From 9e46a79c17d4f0153f8347a17fc18710e10a8298 Mon Sep 17 00:00:00 2001 From: "rongfu.leng" Date: Wed, 15 Apr 2026 08:52:32 +0800 Subject: [PATCH 01/97] [Bugfix] Handle numpy array outputs when generate image (#1680) Signed-off-by: rongfu.leng --- .../openai_api/test_image_server.py | 88 +++++++++++++++++++ vllm_omni/entrypoints/openai/api_server.py | 35 +++++++- 2 files changed, 122 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index c91c5a5c7..4b38692da 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -1165,3 +1165,91 @@ def test_image_edit_with_seed_zero_single_stage(test_client): f"Expected seed=0, but got seed={captured_sampling_params.seed}. " "This indicates the bug where seed=0 is treated as falsy." ) + + +def test_normalize_image(): + """Test _normalize_image with various input types""" + import numpy as np + + from vllm_omni.entrypoints.openai.api_server import _normalize_image + + # Test PIL Image input + img = Image.new("RGB", (64, 64), color="red") + result = _normalize_image(img) + assert isinstance(result, Image.Image) + assert result.size == (64, 64) + + # Test uint8 numpy array + arr = np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8) + result = _normalize_image(arr) + assert isinstance(result, Image.Image) + assert result.size == (64, 64) + + # Test float [0, 1] numpy array + arr = np.random.rand(64, 64, 3).astype(np.float32) + result = _normalize_image(arr) + assert isinstance(result, Image.Image) + assert result.size == (64, 64) + + # Test float [-1, 1] numpy array + arr = np.random.rand(64, 64, 3).astype(np.float32) * 2 - 1 + result = _normalize_image(arr) + assert isinstance(result, Image.Image) + assert result.size == (64, 64) + + # Test batch dimensions (1, 1, H, W, C) + arr = np.random.randint(0, 255, (1, 1, 64, 64, 3), dtype=np.uint8) + result = _normalize_image(arr) + assert isinstance(result, Image.Image) + assert result.size == (64, 64) + + +def test_extract_images_from_result(): + """Test _extract_images_from_result with various result formats""" + import numpy as np + + from vllm_omni.entrypoints.openai.api_server import _extract_images_from_result + + # Test empty result + class EmptyResult: + pass + + result = EmptyResult() + images = _extract_images_from_result(result) + assert images == [] + + # Test nested batch: [np.array(shape=(3, 64, 64, 3))] + batch = np.random.randint(0, 255, (3, 1, 64, 64, 3), dtype=np.uint8) + + class BatchResult: + def __init__(self): + self.images = [batch] + + result = BatchResult() + images = _extract_images_from_result(result) + assert len(images) == 3 + assert all(isinstance(img, Image.Image) for img in images) + assert all(img.size == (64, 64) for img in images) + + # Test dict path: result.request_output["images"] + class DictRequestOutput: + def __init__(self): + self.request_output = {"images": [np.random.randint(0, 255, (64, 64, 3), dtype=np.uint8)]} + + result = DictRequestOutput() + images = _extract_images_from_result(result) + assert len(images) == 1 + assert isinstance(images[0], Image.Image) + + # Test attribute path: result.request_output.images + class AttrRequestOutput: + def __init__(self): + self.request_output = type( + "obj", (), {"images": [np.random.randint(0, 255, (32, 32, 3), dtype=np.uint8)]} + )() + + result = AttrRequestOutput() + images = _extract_images_from_result(result) + assert len(images) == 1 + assert isinstance(images[0], Image.Image) + assert images[0].size == (32, 32) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 6a65f4433..d847a96db 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -18,6 +18,7 @@ from http import HTTPStatus from typing import Annotated, Any, Literal, cast import httpx +import numpy as np import vllm.envs as envs from fastapi import APIRouter, Depends, File, Form, HTTPException, Query, Request, UploadFile, WebSocket from fastapi.responses import FileResponse, JSONResponse, Response, StreamingResponse @@ -1767,6 +1768,34 @@ def _update_if_not_none(object: Any, key: str, val: Any) -> None: setattr(object, key, val) +def _normalize_image(image: Any) -> Any: + """Normalize a single image output to a PIL-compatible format.""" + if isinstance(image, Image.Image): + return image + if not isinstance(image, np.ndarray): + raise ValueError(f"Unsupported image type: {type(image)}") + if not np.issubdtype(image.dtype, np.integer) and not np.issubdtype(image.dtype, np.floating): + raise ValueError(f"Unsupported dtype: {image.dtype}") + if isinstance(image, np.ndarray): + while image.ndim > 3: + image = image[0] + if image.min() < 0: + if image.min() < -1.01 or image.max() > 1.01: + logger.warning( + f"Image float range [{image.min():.2f}, {image.max():.2f}] outside expected [-1, 1]. " + f"Clipping to [-1, 1] before normalization." + ) + image = np.clip(image, -1.0, 1.0) * 0.5 + 0.5 + elif image.max() > 1.01: + logger.warning( + f"Image float range [{image.min():.2f}, {image.max():.2f}] outside expected [0, 1]. " + f"Clipping to [0, 1] before normalization." + ) + image = (np.clip(image, 0.0, 1.0) * 255).astype(np.uint8) + image = Image.fromarray(image) + return image + + def _extract_images_from_result(result: Any) -> list[Any]: images = [] if hasattr(result, "images") and result.images: @@ -1777,6 +1806,10 @@ def _extract_images_from_result(result: Any) -> list[Any]: images = request_output["images"] elif hasattr(request_output, "images") and request_output.images: images = request_output.images + # Handle when generate more than one image + if images and isinstance(images[0], np.ndarray) and images[0].shape[0] > 1 and images[0].ndim == 5: + # Unwrap batch: (N, T, H, W, C) -> [img1, img2, ...] + images = list(images[0]) # Flatten nested lists (e.g., from layered models like Qwen-Image-Layered). # Note: This only flattens one level deep. Deeper nesting is not supported. flattened = [] @@ -1785,7 +1818,7 @@ def _extract_images_from_result(result: Any) -> list[Any]: flattened.extend(img) else: flattened.append(img) - return flattened + return [_normalize_image(img) for img in flattened] async def _load_input_images( From 02e5dc747d028ab75a136988985b32dc83d33557 Mon Sep 17 00:00:00 2001 From: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Date: Tue, 14 Apr 2026 21:58:25 -0400 Subject: [PATCH 02/97] [Perf] VoxCPM2: streaming VAE + compile optimization (45% RTF reduction) (#2758) Signed-off-by: Yueqian Lin --- examples/offline_inference/voxcpm2/end2end.py | 10 +- .../online_serving/voxcpm2/gradio_demo.py | 602 ++++++++++++++++++ tests/e2e/offline_inference/test_voxcpm2.py | 10 +- vllm_omni/engine/output_processor.py | 7 +- .../models/voxcpm2/minicpm4_paged.py | 71 +-- .../models/voxcpm2/voxcpm2_talker.py | 95 ++- 6 files changed, 716 insertions(+), 79 deletions(-) create mode 100644 examples/online_serving/voxcpm2/gradio_demo.py diff --git a/examples/offline_inference/voxcpm2/end2end.py b/examples/offline_inference/voxcpm2/end2end.py index ce404bf96..687e59601 100644 --- a/examples/offline_inference/voxcpm2/end2end.py +++ b/examples/offline_inference/voxcpm2/end2end.py @@ -74,16 +74,20 @@ def extract_audio(multimodal_output: dict) -> torch.Tensor: The output processor concatenates per-step delta tensors under ``model_outputs``. Falls back to ``audio`` for backwards compat. """ - audio = multimodal_output.get("model_outputs") or multimodal_output.get("audio") + audio = multimodal_output.get("model_outputs") + if audio is None: + audio = multimodal_output.get("audio") if audio is None: raise ValueError(f"No audio key in multimodal_output: {list(multimodal_output.keys())}") if isinstance(audio, list): - # Take the last valid tensor (most complete audio) + # Defensive: usually the output processor consolidates into a single + # tensor at request completion, but concatenate here too in case the + # caller consumes intermediate (pre-consolidation) outputs. valid = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio if a is not None] if not valid: raise ValueError("Audio list is empty or all elements are None.") - return valid[-1] + return torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] return torch.as_tensor(audio).float().cpu().reshape(-1) diff --git a/examples/online_serving/voxcpm2/gradio_demo.py b/examples/online_serving/voxcpm2/gradio_demo.py new file mode 100644 index 000000000..a33a2d924 --- /dev/null +++ b/examples/online_serving/voxcpm2/gradio_demo.py @@ -0,0 +1,602 @@ +"""Gradio demo for VoxCPM2 TTS with gapless streaming audio playback. + +Uses a custom AudioWorklet-based player for gap-free streaming +(adapted from the Qwen3-TTS demo). Audio is streamed from the vLLM +server through a same-origin proxy and played via the Web Audio API's +AudioWorklet, which maintains a FIFO buffer queue and plays samples at +the audio clock rate. + +Usage: + # Start the vLLM server first: + python -m vllm_omni.entrypoints.openai.api_server \ + --model openbmb/VoxCPM2 \ + --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm2.yaml \ + --host 0.0.0.0 --port 8000 + + # Then launch the demo: + python gradio_demo.py --api-base http://localhost:8000 +""" + +from __future__ import annotations + +import argparse +import base64 +import io +import json +import logging + +import gradio as gr +import httpx +import numpy as np +import soundfile as sf +from fastapi import FastAPI, Request +from fastapi.responses import Response, StreamingResponse + +logger = logging.getLogger(__name__) + +SAMPLE_RATE = 48000 + +# ── AudioWorklet processor (loaded in browser via Blob URL) ────────── +WORKLET_JS = r""" +class TTSPlaybackProcessor extends AudioWorkletProcessor { + constructor() { + super(); + this.queue = []; + this.buf = null; + this.pos = 0; + this.playing = false; + this.played = 0; + this.port.onmessage = (e) => { + if (e.data && e.data.type === 'clear') { + this.queue = []; this.buf = null; this.pos = 0; this.played = 0; + if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped'}); } + return; + } + this.queue.push(e.data); + }; + } + process(inputs, outputs) { + const out = outputs[0][0]; + for (let i = 0; i < out.length; i++) { + if (!this.buf || this.pos >= this.buf.length) { + if (this.queue.length > 0) { + this.buf = this.queue.shift(); this.pos = 0; + } else { + for (let j = i; j < out.length; j++) out[j] = 0; + if (this.playing) { this.playing = false; this.port.postMessage({type:'stopped', played:this.played}); } + return true; + } + } + out[i] = this.buf[this.pos++] / 32768; + this.played++; + } + if (!this.playing) { this.playing = true; this.port.postMessage({type:'started'}); } + return true; + } +} +registerProcessor('tts-playback-processor', TTSPlaybackProcessor); +""" + +PLAYER_HTML = """ +
+
+
+ Ready + +
+ + + +
+""" + + +def _build_player_js() -> str: + return f""" + +""" + + +def _encode_audio(audio_data: tuple) -> str: + sr, audio_np = audio_data + if audio_np.dtype in (np.float32, np.float64): + audio_np = np.clip(audio_np, -1.0, 1.0) + audio_np = (audio_np * 32767).astype(np.int16) + elif audio_np.dtype != np.int16: + audio_np = audio_np.astype(np.int16) + buf = io.BytesIO() + sf.write(buf, audio_np, sr, format="WAV") + return f"data:audio/wav;base64,{base64.b64encode(buf.getvalue()).decode()}" + + +def create_app(api_base: str): + app = FastAPI() + _pending: dict[str, dict] = {} + + @app.post("/proxy/v1/audio/speech") + async def proxy_speech(request: Request): + body = await request.json() + req_id = body.get("_req_id") + if req_id and req_id in _pending: + body = _pending.pop(req_id) + logger.info("Proxy: %s", {k: (f"<{len(str(v))} chars>" if k == "ref_audio" else v) for k, v in body.items()}) + try: + client = httpx.AsyncClient(timeout=300) + resp = await client.send( + client.build_request( + "POST", + f"{api_base}/v1/audio/speech", + json=body, + headers={"Authorization": "Bearer EMPTY", "Content-Type": "application/json"}, + ), + stream=True, + ) + except Exception as exc: + logger.exception("Proxy connection error") + await client.aclose() + return Response(content=str(exc), status_code=502) + if resp.status_code != 200: + content = await resp.aread() + await resp.aclose() + await client.aclose() + return Response(content=content, status_code=resp.status_code) + + async def relay(): + try: + async for chunk in resp.aiter_bytes(): + yield chunk + finally: + await resp.aclose() + await client.aclose() + + return StreamingResponse(relay(), media_type="application/octet-stream") + + css = """ + #generate-btn button { width: 100%; } + #streaming-player { border: 1px solid var(--border-color-primary) !important; border-radius: var(--block-radius) !important; padding: var(--block-padding) !important; } + """ + theme = gr.themes.Default( + primary_hue=gr.themes.Color( + c50="#f0f5ff", + c100="#dce6f9", + c200="#b8cef3", + c300="#8eb2eb", + c400="#6496e0", + c500="#4A90D9", + c600="#3a7bc8", + c700="#2d66b0", + c800="#1f4f8f", + c900="#163a6e", + c950="#0e2650", + ), + ) + + with gr.Blocks(title="VoxCPM2 TTS Demo") as demo: + gr.HTML(f""" +
+ vLLM-Omni +
+

VoxCPM2 Streaming Demo

+ + Served by vLLM-Omni + · {api_base} + · 48 kHz + +
+
+ """) + + gr.Markdown( + "**Three modes:** " + "**Voice Design** (control instruction only) · " + "**Controllable Cloning** (ref audio + optional style control) · " + "**Ultimate Cloning** (ref audio + transcript for audio continuation)" + ) + + with gr.Row(): + with gr.Column(scale=3): + text_input = gr.Textbox( + label="Target Text", + placeholder="Enter text to synthesize...", + lines=4, + ) + control_instruction = gr.Textbox( + label="Control Instruction (optional)", + placeholder="e.g. A warm young woman / Excited and fast-paced", + lines=2, + info="Describe voice style, emotion, pace. Works for both Voice Design and Controllable Cloning.", + ) + + with gr.Accordion("Voice Cloning", open=False): + ref_audio = gr.Audio( + label="Reference Audio (upload for cloning)", + type="numpy", + sources=["upload", "microphone"], + ) + ref_audio_url = gr.Textbox( + label="or Reference Audio URL", + placeholder="https://example.com/reference.wav", + ) + ultimate_clone = gr.Checkbox( + label="Ultimate Cloning Mode", + value=False, + info="Provide transcript of ref audio for audio continuation (disables control instruction)", + ) + prompt_text = gr.Textbox( + label="Reference Audio Transcript", + placeholder="Transcript of your reference audio (for ultimate cloning)", + lines=2, + visible=False, + ) + + with gr.Row(): + stream_checkbox = gr.Checkbox( + label="Stream (gapless)", + value=True, + info="AudioWorklet streaming", + ) + with gr.Row(): + generate_btn = gr.Button( + "Generate Speech", + variant="primary", + size="lg", + elem_id="generate-btn", + scale=3, + ) + reset_btn = gr.Button("Reset", variant="secondary", size="lg", scale=1) + + with gr.Column(scale=2): + player_html = gr.HTML( + value=PLAYER_HTML, + visible=True, + label="streaming player", + elem_id="streaming-player", + ) + audio_output = gr.Audio( + label="generated audio", + interactive=False, + autoplay=True, + visible=False, + ) + gr.Examples( + examples=[ + ["Hello, this is a VoxCPM2 demo running on vLLM-Omni.", ""], + [ + "I have a dream that my four little children will one day live in a nation " + "where they will not be judged by the color of their skin but by the content " + "of their character.", + "", + ], + [ + "I never asked you to stay. It's not like I care or anything. " + "But why does it still hurt so much now that you're gone?", + "A young girl with a soft, sweet voice. Speaks slowly with a melancholic tone.", + ], + ], + inputs=[text_input, control_instruction], + label="examples", + ) + gr.HTML(""" +
+ + vLLM-Omni + +
+ """) + + hidden_payload = gr.Textbox(visible=False, elem_id="tts-payload") + + def on_ultimate_toggle(checked): + return ( + gr.update(visible=checked), # prompt_text + gr.update(interactive=not checked), # control_instruction + ) + + ultimate_clone.change( + fn=on_ultimate_toggle, + inputs=[ultimate_clone], + outputs=[prompt_text, control_instruction], + ) + + def on_stream_change(stream: bool): + if stream: + return gr.update(visible=True), gr.update(visible=False) + return gr.update(visible=False), gr.update(visible=True) + + stream_checkbox.change( + fn=on_stream_change, + inputs=[stream_checkbox], + outputs=[player_html, audio_output], + ) + + def on_reset(): + return "", "", None, "", False, "", PLAYER_HTML + + reset_btn.click( + fn=on_reset, + outputs=[ + text_input, + control_instruction, + audio_output, + hidden_payload, + ultimate_clone, + prompt_text, + player_html, + ], + js="() => { if (window.ttsStop) window.ttsStop(); }", + ) + + def on_generate(stream_enabled, text, ctrl_instr, ref_a, ref_url, ult_clone, p_text): + import time as _time + + if not text or not text.strip(): + raise gr.Error("Please enter text to synthesize.") + + # VoxCPM2 uses "(instruction)text" format for control + ctrl = ctrl_instr.strip() if ctrl_instr and not ult_clone else "" + final_text = f"({ctrl}){text.strip()}" if ctrl else text.strip() + + payload: dict = { + "input": final_text, + "voice": "default", + "response_format": "pcm" if stream_enabled else "wav", + "stream": stream_enabled, + } + + # Reference audio for cloning + ref_url_s = ref_url.strip() if ref_url else "" + if ref_url_s: + payload["ref_audio"] = ref_url_s + elif ref_a is not None: + payload["ref_audio"] = _encode_audio(ref_a) + + # Ultimate cloning: prompt_audio + prompt_text for continuation + if ult_clone and p_text and p_text.strip(): + if ref_url_s: + payload["prompt_audio"] = ref_url_s + elif ref_a is not None: + payload["prompt_audio"] = payload.get("ref_audio", "") + payload["prompt_text"] = p_text.strip() + + if stream_enabled: + if ref_a is not None and not ref_url_s: + req_id = f"req-{int(_time.time() * 1000)}" + _pending[req_id] = payload + browser_payload = {"_req_id": req_id, "_nonce": int(_time.time() * 1000)} + return json.dumps(browser_payload), gr.update() + payload["_nonce"] = int(_time.time() * 1000) + return json.dumps(payload), gr.update() + else: + try: + with httpx.Client(timeout=300.0) as client: + resp = client.post( + f"{api_base}/v1/audio/speech", + json=payload, + headers={"Content-Type": "application/json", "Authorization": "Bearer EMPTY"}, + ) + except httpx.ConnectError: + raise gr.Error(f"Cannot connect to server at {api_base}.") + if resp.status_code != 200: + raise gr.Error(f"Server error ({resp.status_code}): {resp.text[:200]}") + audio_np, sr = sf.read(io.BytesIO(resp.content)) + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + return "", (sr, audio_np.astype(np.float32)) + + generate_btn.click( + fn=on_generate, + inputs=[ + stream_checkbox, + text_input, + control_instruction, + ref_audio, + ref_audio_url, + ultimate_clone, + prompt_text, + ], + outputs=[hidden_payload, audio_output], + ).then( + fn=lambda p: p, + inputs=[hidden_payload], + outputs=[hidden_payload], + js="(p) => { if (p && p.trim()) { const d = JSON.parse(p); delete d._nonce; window.ttsGenerate(d); } return p; }", + ) + + demo.queue() + + return gr.mount_gradio_app(app, demo, path="/", css=css, theme=theme, head=_build_player_js()) + + +def main(): + parser = argparse.ArgumentParser(description="VoxCPM2 streaming Gradio demo") + parser.add_argument("--api-base", default="http://localhost:8000", help="vLLM API server URL") + parser.add_argument("--host", default="0.0.0.0", help="Gradio server host") + parser.add_argument("--port", type=int, default=7860, help="Gradio server port") + args = parser.parse_args() + + logging.basicConfig(level=logging.INFO) + print(f"Connecting to vLLM server at: {args.api_base}") + + import uvicorn + + uvicorn.run(create_app(args.api_base), host=args.host, port=args.port) + + +if __name__ == "__main__": + main() diff --git a/tests/e2e/offline_inference/test_voxcpm2.py b/tests/e2e/offline_inference/test_voxcpm2.py index 4e4f635d5..6ec4630a4 100644 --- a/tests/e2e/offline_inference/test_voxcpm2.py +++ b/tests/e2e/offline_inference/test_voxcpm2.py @@ -33,14 +33,16 @@ def _extract_audio(multimodal_output: dict) -> torch.Tensor: """Extract the final complete audio tensor from multimodal output.""" assert isinstance(multimodal_output, dict), f"Expected dict, got {type(multimodal_output)}" - # Output processor accumulates per-step full audio under "audio". - audio = multimodal_output.get("audio") or multimodal_output.get("model_outputs") + # Output processor accumulates per-step audio chunks under "audio". + audio = multimodal_output.get("audio") + if audio is None: + audio = multimodal_output.get("model_outputs") assert audio is not None, f"No audio key, got {list(multimodal_output.keys())}" if isinstance(audio, list): - valid = [x for x in audio if isinstance(x, torch.Tensor) and x.numel() > 100] + valid = [torch.as_tensor(x).float().cpu().reshape(-1) for x in audio if x is not None] assert valid, "No valid audio tensors in output list" - audio = valid[-1] + audio = torch.cat(valid, dim=0) if len(valid) > 1 else valid[0] assert isinstance(audio, torch.Tensor), f"Expected Tensor, got {type(audio)}" return audio diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 43d02e85b..badd799fc 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -118,9 +118,10 @@ class OmniRequestState(RequestState): if isinstance(v, list) and v and isinstance(v[0], torch.Tensor): try: if k == "audio": - # When the audio tensor shape is inconsistent, torch.cat will fail. - # We need to use torch.cat in -1 dimension. - continue + # Concatenate delta audio chunks (1-D) into the full waveform. + # Each entry is a per-step slice; flatten to -1 so chunks with + # inconsistent leading dims can still be joined on the sample axis. + self.mm_accumulated[k] = torch.cat([t.reshape(-1) for t in v], dim=0) elif k == "sr": # Sample rate is a constant scalar, keep last value. self.mm_accumulated[k] = v[-1] diff --git a/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py b/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py index 7ea5bc229..40bacfff6 100644 --- a/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py +++ b/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py @@ -308,31 +308,28 @@ class MiniCPM4PagedForVoxCPM2(nn.Module): return hidden_states def compile_selective(self) -> list[str]: - """Compile MLP + o_proj; keep RMSNorm/RoPE eager for precision.""" - compiled: list[str] = [] - for i, layer in enumerate(self.layers): - if i in self._compiled_layers: - continue - try: - layer.mlp = torch.compile( - layer.mlp, - mode="default", - fullgraph=True, - ) - layer.self_attn.o_proj = torch.compile( - layer.self_attn.o_proj, - mode="default", - fullgraph=True, - ) - layer.self_attn._fused_qkv_weight = None - self._compiled_layers.add(i) - if i == 0: - compiled.append(f"layers.*.mlp (×{len(self.layers)})") - compiled.append(f"layers.*.self_attn.o_proj (×{len(self.layers)})") - except Exception as e: - logger.warning("compile_selective: layer %d failed: %s", i, e) - break - return compiled + """Compile the full model forward as one graph. + + Earlier versions compiled ``layer.mlp`` + ``layer.self_attn.o_proj`` + (PR #2690) and then the whole ``layer`` (perf/voxcpm2-streaming-vae). + Both still paid one Dynamo dispatch per layer per decode step. + V3 profiling showed 1,332 per-layer dispatches (~28 layers × ~47 + decode steps) costing ~726 ms of CPU self-time for a long prompt. + + Compiling ``forward`` at the model level lets Dynamo unroll the + 28-layer Python loop inside the graph. Graph breaks at + PagedAttention produce sub-graphs but Dynamo memoises the whole + trace once, so the per-step dispatch drops from 28 to just a few. + """ + if self._compiled_layers: + return [] + # Null the fused-qkv caches so the compile sees the real weight layout. + for layer in self.layers: + layer.self_attn._fused_qkv_weight = None + self.forward = torch.compile(self.forward, mode="default", fullgraph=False) + # Mark every layer as compiled so idempotent callers don't double-wrap. + self._compiled_layers.update(range(len(self.layers))) + return ["forward (whole model)"] def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights from native checkpoint (base_lm. prefix pre-stripped).""" @@ -415,22 +412,14 @@ class MiniCPM4PagedResidualLM(nn.Module): return hidden_states def compile_selective(self) -> list[str]: - """Compile MLP + o_proj (same as base_lm).""" - compiled: list[str] = [] - for i, layer in enumerate(self.layers): - if i in self._compiled_layers: - continue - try: - layer.mlp = torch.compile(layer.mlp, mode="default", fullgraph=True) - layer.self_attn.o_proj = torch.compile(layer.self_attn.o_proj, mode="default", fullgraph=True) - layer.self_attn._fused_qkv_weight = None - self._compiled_layers.add(i) - if i == 0: - compiled.append(f"layers.*.mlp (×{len(self.layers)})") - compiled.append(f"layers.*.self_attn.o_proj (×{len(self.layers)})") - except Exception as e: - logger.warning("compile_selective: residual layer %d failed: %s", i, e) - return compiled + """Compile the full residual model forward as one graph (same strategy as base_lm).""" + if self._compiled_layers: + return [] + for layer in self.layers: + layer.self_attn._fused_qkv_weight = None + self.forward = torch.compile(self.forward, mode="default", fullgraph=False) + self._compiled_layers.update(range(len(self.layers))) + return ["forward (whole residual)"] def load_weights_from_native(self, native_residual_lm: nn.Module) -> int: """Load weights from native residual_lm. Returns param count.""" diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 0898ca59a..94f065890 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -11,6 +11,7 @@ Architecture: from __future__ import annotations import dataclasses +import logging import os import time from collections.abc import Iterable @@ -19,7 +20,6 @@ from typing import Any import librosa import torch import torch.nn as nn -from einops import rearrange from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.utils import ( @@ -86,7 +86,11 @@ class _RequestState: curr_prefix_feat_cond: torch.Tensor | None = None last_audio_patch_gpu: torch.Tensor | None = None precomputed_stop_logits: torch.Tensor | None = None - accumulated_patches: list[torch.Tensor] = dataclasses.field(default_factory=list) + # Rolling tail of previously-decoded latents used as VAE receptive-field context. + # Shape (n_pad_frames, feat_dim) on GPU. None before first decode. + decode_pad: torch.Tensor | None = None + # Audio chunks already emitted (CPU float32), concatenated for cumulative output. + audio_chunks: list[torch.Tensor] = dataclasses.field(default_factory=list) decode_step_count: int = 0 request_start_time: float = 0.0 prefill_completed: bool = False @@ -229,11 +233,11 @@ def _optimized_solve_euler( buffers.x_in[b : 2 * b].copy_(x) buffers.mu_in[:b].copy_(mu) buffers.mu_in[b : 2 * b].zero_() - buffers.t_in[:b].fill_(t.item()) - buffers.t_in[b : 2 * b].fill_(t.item()) + # Broadcast the 0-dim GPU scalar directly instead of + # ``.fill_(t.item())`` — ``.item()`` forces a GPU->CPU sync. + buffers.t_in[: 2 * b].copy_(t) if mean_mode: - buffers.dt_in[:b].fill_(dt.item()) - buffers.dt_in[b : 2 * b].fill_(dt.item()) + buffers.dt_in[: 2 * b].copy_(dt) else: buffers.dt_in.zero_() buffers.cond_in[:b].copy_(cond[:b]) @@ -263,9 +267,10 @@ def _optimized_solve_euler( else: buffers.x_in[:b].copy_(x) buffers.mu_in[:b].copy_(mu) - buffers.t_in[:b].fill_(t.item()) + # Broadcast the 0-dim GPU scalar; ``.fill_(t.item())`` would sync. + buffers.t_in[:b].copy_(t) if mean_mode: - buffers.dt_in[:b].fill_(dt.item()) + buffers.dt_in[:b].copy_(dt) else: buffers.dt_in[:b].zero_() buffers.cond_in[:b].copy_(cond[:b]) @@ -320,7 +325,10 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._inference_timesteps = 10 self._cfg_value = 2.0 self._cfg_cutoff_ratio = 1.0 - self._vae_decode_interval = 5 + # Number of trailing latent frames to keep as VAE receptive-field context + # for sliding-window streaming decode. 12 matches the nanovllm reference + # implementation and covers the longest VAE decoder receptive field. + self._n_decode_pad_frames = 12 self._enable_torch_compile = True self._compile_vae = True self._max_decode_steps = 2000 @@ -686,7 +694,9 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): state.request_start_time = time.perf_counter() state.prefill_completed = True - logger.info("PREFILL[%s]: patch norm=%.4f", state.request_id, pred_feat.norm().item()) + if logger.isEnabledFor(logging.DEBUG): + # Only compute the norm (which forces a GPU->CPU sync) if we will log it. + logger.debug("PREFILL[%s]: patch norm=%.4f", state.request_id, pred_feat.norm().item()) self._perf.reset() def _finish_decode(self, state: _RequestState, meta: dict, res_out: torch.Tensor, dev: Any): @@ -720,26 +730,54 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): # -------------------- audio collection -------------------- def _collect_audio(self, state: _RequestState) -> torch.Tensor | None: + """Per-step sliding-window VAE decode (nanovllm pattern). + + Each decode step feeds ``[decode_pad, new_patch]`` through the VAE + and slices out only the audio region corresponding to the new patch. + The pad buffer (last ``_n_decode_pad_frames`` latent frames) provides + the receptive-field context needed by the VAE's transposed convolutions, + eliminating boundary artifacts between chunks. + + Returns the delta audio chunk (not cumulative) so the output processor + can stream each chunk to the client independently. + """ patch = state.last_audio_patch_gpu - if patch is not None: - state.last_audio_patch_gpu = None - state.accumulated_patches.append(patch.reshape(1, -1).float()) - - if not state.accumulated_patches: + if patch is None: return None + state.last_audio_patch_gpu = None - n = len(state.accumulated_patches) - if n <= 1 or n % self._vae_decode_interval == 0 or state.is_stopping: - self._perf.start("vae_decode") - all_p = torch.cat(state.accumulated_patches, dim=0) - state.accumulated_patches = [all_p] - feat = rearrange(all_p.reshape(1, -1, self._feat_dim), "b t d -> b d t") - with torch.no_grad(): - audio = self.tts.audio_vae.decode(feat.to(self._device)).reshape(-1).cpu().float() - self._perf.stop("vae_decode") - state.last_decoded_audio = audio - return audio - return state.last_decoded_audio + # patch shape: (patch_size, feat_dim) or (1, patch_size, feat_dim) + new_latent = patch.reshape(-1, self._feat_dim).to(torch.float32) + n_new = new_latent.shape[0] # = patch_size (typically 4) + + self._perf.start("vae_decode") + + # Build VAE input: [pad_frames | new_latent] + if state.decode_pad is not None: + vae_input = torch.cat([state.decode_pad, new_latent], dim=0) + pad_frames = state.decode_pad.shape[0] + else: + vae_input = new_latent + pad_frames = 0 + + # VAE decode: (1, feat_dim, T_frames) -> (1, 1, T_samples) + feat = vae_input.unsqueeze(0).transpose(1, 2).contiguous() + with torch.no_grad(): + audio = self.tts.audio_vae.decode(feat.to(self._device)).reshape(-1) + + # Slice out only the new audio (after the pad region). + # Each latent frame maps to decoder_chunk_size audio samples. + dcs = int(getattr(self.tts.audio_vae, "decode_chunk_size", audio.numel() // vae_input.shape[0])) + new_audio = audio[pad_frames * dcs : (pad_frames + n_new) * dcs].detach().cpu().float() + + # Roll the pad buffer: keep last N latent frames as context for next step. + all_latents = vae_input # [pad + new] + state.decode_pad = all_latents[-self._n_decode_pad_frames :].detach() + + state.audio_chunks.append(new_audio) + state.last_decoded_audio = new_audio + self._perf.stop("vae_decode") + return new_audio # -------------------- compute_logits -------------------- @@ -830,7 +868,8 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): state = self._get_or_create_state(req_id) state.prefill_text = "" - state.accumulated_patches = [] + state.decode_pad = None + state.audio_chunks = [] state.prefill_completed = False state.decode_step_count = 0 state.precomputed_stop_logits = None From a782ae47805d9761f446e4e715530af0f54859ab Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Wed, 15 Apr 2026 10:28:28 +0800 Subject: [PATCH 03/97] [Perf] Enhance benchmark script to support baseline thresholds and proved result handling (#2789) --- tests/dfx/perf/scripts/run_benchmark.py | 98 +++++++++++++++++-- .../scripts/test_benchmark_stability.py | 2 + tools/nightly/generate_nightly_perf_excel.py | 49 +++++++--- tools/nightly/generate_nightly_perf_html.py | 45 +++++++-- 4 files changed, 162 insertions(+), 32 deletions(-) diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index b64cc0d95..67dedcd04 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -56,16 +56,41 @@ def omni_server(request): print("OmniServer stopped") +def _safe_filename_token(value: Any | None, *, default: str = "na") -> str: + """Make a single path segment safe for result filenames on common filesystems.""" + if value is None: + return default + s = str(value).strip() + for bad in ("/", "\\", ":", "*", "?", '"', "<", ">", "|"): + s = s.replace(bad, "_") + return s if s else default + + def run_benchmark( args: list, test_name: str, flow, dataset_name: str, num_prompt, + *, + 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, ) -> Any: - """Run a single benchmark iteration and return the parsed result JSON.""" + """Run a single benchmark iteration and return the parsed result JSON. + + After ``vllm bench`` writes the JSON, ``result["baseline"]`` holds the same + per-metric resolved thresholds as ``assert_result`` (via ``_baseline_thresholds_for_step``). + When ``random_input_len`` / ``random_output_len`` are set, they are also written into the result JSON; + omitted keys when not configured. + """ current_dt = datetime.now().strftime("%Y%m%d-%H%M%S") - result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_{current_dt}.json" + ri = _safe_filename_token(random_input_len) + ro = _safe_filename_token(random_output_len) + result_filename = f"result_{test_name}_{dataset_name}_{flow}_{num_prompt}_in{ri}_out{ro}_{current_dt}.json" if "--result-filename" in args: print(f"The result file will be overwritten by {result_filename}") command = ( @@ -97,8 +122,26 @@ def run_benchmark( else: result_dir = "./" - with open(os.path.join(result_dir, result_filename), encoding="utf-8") as f: + result_path = os.path.join(result_dir, result_filename) + with open(result_path, encoding="utf-8") as f: 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, + ) + else: + result["baseline"] = {} + if random_input_len is not None: + result["random_input_len"] = random_input_len + if random_output_len is not None: + result["random_output_len"] = random_output_len + with open(result_path, "w", encoding="utf-8") as f: + json.dump(result, f, ensure_ascii=False, indent=2) + return result @@ -164,10 +207,33 @@ def _resolve_baseline_value( 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 ``test.json`` ``baseline`` block to one threshold per metric (same as ``assert_result``).""" + return { + metric_name: _resolve_baseline_value( + baseline_raw, + sweep_index=sweep_index, + max_concurrency=max_concurrency, + request_rate=request_rate, + ) + for metric_name, baseline_raw in baseline_data.items() + } + + def assert_result( result, params, @@ -179,14 +245,14 @@ def assert_result( ) -> None: assert result["completed"] == num_prompt, "Request failures exist" baseline_data = params.get("baseline", {}) - for metric_name, baseline_raw in baseline_data.items(): + thresholds = _baseline_thresholds_for_step( + baseline_data, + sweep_index=sweep_index, + max_concurrency=max_concurrency, + request_rate=request_rate, + ) + for metric_name, baseline_value in thresholds.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( @@ -258,6 +324,12 @@ def test_performance_benchmark(omni_server, benchmark_params): flow=qps, 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"), ) assert_result( result, @@ -276,6 +348,12 @@ def test_performance_benchmark(omni_server, benchmark_params): flow=concurrency, 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"), ) assert_result( result, diff --git a/tests/dfx/stability/scripts/test_benchmark_stability.py b/tests/dfx/stability/scripts/test_benchmark_stability.py index e8568652d..a9faae8ab 100644 --- a/tests/dfx/stability/scripts/test_benchmark_stability.py +++ b/tests/dfx/stability/scripts/test_benchmark_stability.py @@ -112,6 +112,8 @@ def _run_one_benchmark_batch( flow=flow, dataset_name=dataset_name, num_prompt=num_prompts, + random_input_len=params.get("random_input_len"), + random_output_len=params.get("random_output_len"), ) return result except (FileNotFoundError, OSError) as e: diff --git a/tools/nightly/generate_nightly_perf_excel.py b/tools/nightly/generate_nightly_perf_excel.py index 5f9eb428b..4bb778531 100644 --- a/tools/nightly/generate_nightly_perf_excel.py +++ b/tools/nightly/generate_nightly_perf_excel.py @@ -319,10 +319,10 @@ def _load_json_file(path: str) -> dict[str, Any] | list[Any] | None: def _parse_from_filename(filename: str) -> dict[str, Any]: - """Parse test-related metadata from a result JSON filename. + """Parse test-related metadata from a ``result_test_*.json`` filename. - Expected pattern (after prefix/suffix stripped): - ____ + Matches ``tests/dfx/perf/scripts/run_benchmark.py`` naming, including optional + ``_in{X}_out{Y}_`` before the timestamp (``na`` when unset). """ name, ext = os.path.splitext(filename) if ext != ".json" or not name.startswith(_RESULT_JSON_PREFIX): @@ -331,22 +331,42 @@ def _parse_from_filename(filename: str) -> dict[str, Any]: core = name[len(_RESULT_JSON_PREFIX) :] parts = core.split("_") if len(parts) < 5: - LOGGER.warning("filename '%s' does not match expected pattern, skip parsing test metadata", filename) + LOGGER.warning( + "filename '%s' does not match expected pattern (need >= 5 segments), skip parsing", + filename, + ) return {} - timestamp = parts[-1] - num_prompts_str = parts[-2] - max_concurrency_str = parts[-3] - dataset_name = parts[-4] - test_name = "_".join(parts[:-4]) if parts[:-4] else "" + idx = len(parts) - 1 + timestamp = parts[idx] + idx -= 1 parsed: dict[str, Any] = {} - if len(timestamp) >= 15: parsed["date"] = timestamp - if dataset_name in DATASET_NAME_ALLOWED: - parsed["dataset_name"] = dataset_name + if idx >= 0 and parts[idx].startswith("out"): + parsed["random_output_len"] = parts[idx][3:] + idx -= 1 + if idx >= 0 and parts[idx].startswith("in"): + parsed["random_input_len"] = parts[idx][2:] + idx -= 1 + + if idx < 3: + LOGGER.warning( + "filename '%s' has too few segments after timestamp / optional in-out (idx=%s)", + filename, + idx, + ) + return parsed + + num_prompts_str = parts[idx] + idx -= 1 + flow_str = parts[idx] + idx -= 1 + dataset_name = parts[idx] + idx -= 1 + test_name = "_".join(parts[: idx + 1]) if idx >= 0 else "" try: parsed["num_prompts"] = int(num_prompts_str) @@ -354,13 +374,16 @@ def _parse_from_filename(filename: str) -> dict[str, Any]: pass try: - parsed["max_concurrency"] = int(max_concurrency_str) + parsed["max_concurrency"] = int(flow_str) except (TypeError, ValueError): pass if test_name: parsed["test_name"] = test_name + if dataset_name in DATASET_NAME_ALLOWED: + parsed["dataset_name"] = dataset_name + return parsed diff --git a/tools/nightly/generate_nightly_perf_html.py b/tools/nightly/generate_nightly_perf_html.py index 05dc48d71..dd5ece790 100644 --- a/tools/nightly/generate_nightly_perf_html.py +++ b/tools/nightly/generate_nightly_perf_html.py @@ -67,6 +67,7 @@ def _load_json_file(path: str) -> dict[str, Any] | None: def _parse_from_filename(filename: str) -> dict[str, Any]: + """Parse ``result_test_*.json`` filenames; same rules as ``generate_nightly_perf_excel``.""" name, ext = os.path.splitext(filename) if ext != ".json" or not name.startswith(_RESULT_JSON_PREFIX): return {} @@ -75,32 +76,58 @@ def _parse_from_filename(filename: str) -> dict[str, Any]: parts = core.split("_") if len(parts) < 5: LOGGER.warning( - "filename '%s' does not match expected pattern, skip parsing test metadata", + "filename '%s' does not match expected pattern (need >= 5 segments), skip parsing", filename, ) return {} - timestamp = parts[-1] - num_prompts_str = parts[-2] - max_concurrency_str = parts[-3] - dataset_name = parts[-4] - test_name = "_".join(parts[:-4]) if parts[:-4] else "" + idx = len(parts) - 1 + timestamp = parts[idx] + idx -= 1 parsed: dict[str, Any] = {} if len(timestamp) >= 15: parsed["date"] = timestamp - if dataset_name in ("random", "random-mm"): - parsed["dataset_name"] = dataset_name + + if idx >= 0 and parts[idx].startswith("out"): + parsed["random_output_len"] = parts[idx][3:] + idx -= 1 + if idx >= 0 and parts[idx].startswith("in"): + parsed["random_input_len"] = parts[idx][2:] + idx -= 1 + + if idx < 3: + LOGGER.warning( + "filename '%s' has too few segments after timestamp / optional in-out (idx=%s)", + filename, + idx, + ) + return parsed + + num_prompts_str = parts[idx] + idx -= 1 + flow_str = parts[idx] + idx -= 1 + dataset_name = parts[idx] + idx -= 1 + test_name = "_".join(parts[: idx + 1]) if idx >= 0 else "" + try: parsed["num_prompts"] = int(num_prompts_str) except (TypeError, ValueError): pass + try: - parsed["max_concurrency"] = int(max_concurrency_str) + parsed["max_concurrency"] = int(flow_str) except (TypeError, ValueError): pass + if test_name: parsed["test_name"] = test_name + + if dataset_name in ("random", "random-mm"): + parsed["dataset_name"] = dataset_name + return parsed From 227bab3038a10ba1bde4c2c9154be428b496a7e3 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Wed, 15 Apr 2026 11:07:35 +0800 Subject: [PATCH 04/97] [Benchmark]Omni-modality model accuracy benchmark(Daily-Omni & seed-tts-eval) (#2558) Signed-off-by: amy-why-3459 --- pyproject.toml | 11 + .../data_modules/daily_omni_dataset.py | 887 ++++++++++++++++++ .../data_modules/daily_omni_eval.py | 406 ++++++++ .../data_modules/daily_omni_text_audio.py | 255 +++++ .../data_modules/seed_tts_dataset.py | 272 ++++++ .../benchmarks/data_modules/seed_tts_eval.py | 729 ++++++++++++++ vllm_omni/benchmarks/patch/__init__.py | 3 + vllm_omni/benchmarks/patch/patch.py | 332 ++++++- vllm_omni/benchmarks/serve.py | 12 + vllm_omni/entrypoints/cli/benchmark/serve.py | 143 ++- 10 files changed, 3041 insertions(+), 9 deletions(-) create mode 100644 vllm_omni/benchmarks/data_modules/daily_omni_dataset.py create mode 100644 vllm_omni/benchmarks/data_modules/daily_omni_eval.py create mode 100644 vllm_omni/benchmarks/data_modules/daily_omni_text_audio.py create mode 100644 vllm_omni/benchmarks/data_modules/seed_tts_dataset.py create mode 100644 vllm_omni/benchmarks/data_modules/seed_tts_eval.py diff --git a/pyproject.toml b/pyproject.toml index 57a4b474f..753e0e398 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -61,6 +61,17 @@ demo = [ "gradio>=6.7.0", ] +# Seed-TTS serve benchmark WER (BytedanceSpeech/seed-tts-eval run_wer.py protocol). +seed-tts-eval = [ + "jiwer>=3.0.0", + "zhon>=2.0.0", + "zhconv>=1.4.2", + "scipy>=1.10.0", + "soundfile>=0.12.0", + "transformers>=4.36.0", + "funasr>=1.0.0", +] + docs = [ "mkdocs>=1.5.0", "mkdocs-api-autonav", diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py new file mode 100644 index 000000000..01b86d0fd --- /dev/null +++ b/vllm_omni/benchmarks/data_modules/daily_omni_dataset.py @@ -0,0 +1,887 @@ +"""Daily-Omni Dataset loader for benchmark. + +Daily-Omni is an audio-visual reasoning benchmark with 684 videos +and 1,197 multiple-choice QA pairs across 6 major task types. + +Dataset source: https://huggingface.co/datasets/liarliar/Daily-Omni + +Supports loading QA metadata from: +- Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments +- HuggingFace datasets (``dataset_path``): legacy online mode + +The videos must be separately downloaded and extracted from Videos.tar. + +Why ``BenchmarkDataset`` instead of ``HuggingFaceDataset``? + vLLM's ``HuggingFaceDataset`` is a thin wrapper whose ``__init__`` always ends by calling + ``load_data()`` → ``datasets.load_dataset(...)`` with a required Hub id and split. That + contract fits "Hub-only" benches, but Daily-Omni also needs **offline QA metadata** from a + local ``qa.json`` without touching the network. Subclassing ``HuggingFaceDataset`` would + mean fighting the parent constructor (fake ``dataset_path``, reordering ``load_data``, or + duplicating half the parent) and would still imply ``datasets`` is always relevant. + + This class therefore inherits only ``BenchmarkDataset`` (minimal: ``dataset_path``, + ``random_seed``, ``self.data``) and implements **two explicit loaders**: + ``_load_from_local_json`` (default path for air-gapped runs) and ``_load_from_huggingface`` + (optional legacy path for users who prefer ``datasets`` + Hub cache). The latter is **not** + inheritance; it is the same Hub rows as before, factored into a helper so one class can + serve both deployment modes without mandatory ``datasets`` when using ``qa_json_path``. + +Usage: + from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset + + # Local JSON mode (recommended) + dataset = DailyOmniDataset( + qa_json_path="/path/to/qa.json", + video_dir="/path/to/Videos", + random_seed=42, + ) + + # HuggingFace mode (legacy, requires network) + dataset = DailyOmniDataset( + dataset_path="liarliar/Daily-Omni", + dataset_split="train", + random_seed=42, + ) + requests = dataset.sample( + tokenizer=tokenizer, + num_requests=100, + output_len=256, + ) +""" + +import base64 +import json +import logging +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +try: + from vllm.benchmarks.datasets import BenchmarkDataset, SampleRequest +except ImportError: + # Fallback: if BenchmarkDataset not available, use base class from same module + from vllm.benchmarks.datasets import HuggingFaceDataset as BenchmarkDataset + from vllm.benchmarks.datasets import SampleRequest +from vllm.tokenizers import TokenizerLike +from vllm.tokenizers.hf import get_cached_tokenizer + +try: + from datasets import load_dataset +except ImportError: + load_dataset = None + +logger = logging.getLogger(__name__) + + +class _ListDatasetIterator: + """Simple iterator wrapper around a list to mimic HuggingFace streaming dataset behavior.""" + + def __init__(self, data: list[dict[str, Any]]) -> None: + self._data = data + self._index = 0 + + def __iter__(self): + self._index = 0 + return self + + def __next__(self) -> dict[str, Any]: + if self._index >= len(self._data): + raise StopIteration + item = self._data[self._index] + self._index += 1 + return item + + def __len__(self) -> int: + return len(self._data) + + def __getitem__(self, idx: int | slice) -> dict[str, Any] | list[dict[str, Any]]: + return self._data[idx] + + +# Aligns with Lliar-liar/Daily-Omni CLI ``--input_mode`` (test_model/*/testmodel.py). +DailyOmniInputMode = Literal["all", "visual", "audio"] + +# ``build_conversation()`` in Daily-Omni ``test_model/Qwen2.5-Omni/testmodel.py`` (verbatim). +DAILY_OMNI_SYSTEM_TEXT = ( + "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, " + "capable of perceiving auditory and visual inputs, as well as generating text and speech." +) + + +@dataclass +class DailyOmniSampleRequest(SampleRequest): + """``SampleRequest`` with Daily-Omni gold labels for post-run accuracy scoring.""" + + daily_omni_gold_answer: str = "" + daily_omni_video_id: str = "" + daily_omni_task_type: str = "" + #: Official qa.json ``video_duration`` (e.g. ``30s``, ``60s``) for leaderboard-style breakdown. + daily_omni_video_duration: str = "" + #: Official ``video_category`` (YouTube-style category string) for per-category accuracy. + daily_omni_video_category: str = "" + #: Extra JSON fields merged into chat-completions ``extra_body`` (e.g. ``mm_processor_kwargs``). + omni_extra_body: dict[str, Any] | None = None + #: Full OpenAI ``messages`` (system + user) mirroring upstream Daily-Omni conversation. + omni_chat_messages: list[dict[str, Any]] | None = None + #: Used only when ``omni_chat_messages`` is None (non-Daily-Omni-style requests). + omni_chat_mm_position: Literal["first", "last"] = "last" + + +class DailyOmniDataset(BenchmarkDataset): + """Daily-Omni audio-visual QA dataset for benchmarking. + + Inherits ``BenchmarkDataset`` only (not ``HuggingFaceDataset``): see module docstring for why + Hub loading lives in ``_load_from_huggingface`` instead of subclassing the HF base class. + + The dataset includes: + - 684 videos from daily life scenarios (available in Videos.tar) + - 1,197 multiple-choice QA pairs in qa.json + - 6 major task categories + + QA metadata can be loaded from: + - Local JSON file (``qa_json_path``): recommended for offline/air-gapped environments + - HuggingFace datasets (``dataset_path``): legacy online mode + + The videos must be separately downloaded and extracted from Videos.tar. + + Args: + qa_json_path: Path to local qa.json file (offline mode, preferred). When provided, + ``dataset_path`` and ``dataset_split`` are ignored. + dataset_path: HuggingFace dataset path (e.g., "liarliar/Daily-Omni"). Used only if + ``qa_json_path`` is not provided (legacy online mode). + dataset_split: Dataset split to use (default: "train"). Used only in online mode. + random_seed: Random seed for shuffling + video_dir: Directory containing extracted video files (default: None) + input_mode: Which modalities to send, matching upstream Daily-Omni ``--input_mode``: + ``all`` — video + WAV (default; official audio-visual protocol); + ``visual`` — video only; + ``audio`` — extracted WAV only (requires ``{video_id}/{video_id}_audio.wav`` under ``video_dir``). + max_duration_seconds: Reserved for future ffprobe-based filtering; currently **not applied** + when building requests (metadata ``video_duration`` is still passed through for eval). + dataset_subset: Optional HuggingFace subset name (``load_dataset(..., name=...)``); used by bench + ``--hf-subset`` / patch. + no_stream: If True, load the Hub split non-streaming (matches bench ``--no-stream``). + inline_local_video: If True, embed local MP4 as ``data:video/mp4;base64,...`` in requests so + the API server does not need ``--allowed-local-media-path`` (large JSON; use for small runs). + When ``input_mode`` is ``audio`` or ``all``, local WAV is embedded the same way + (``data:audio/wav;base64,...``). + trust_remote_code: Whether to trust remote code when loading HuggingFace dataset + (online mode only). + """ + + SUPPORTED_DATASET_PATHS: set[str] = { + "liarliar/Daily-Omni", + } + #: Default Hub id for synthetic video URLs when ``qa_json_path`` is used (``dataset_path`` None). + DEFAULT_HF_DATASET_ID = "liarliar/Daily-Omni" + IS_MULTIMODAL = True + DEFAULT_OUTPUT_LEN = 256 + + def __init__( + self, + qa_json_path: str | None = None, + dataset_path: str | None = None, + dataset_split: str = "train", + random_seed: int = 0, + video_dir: str | None = None, + input_mode: DailyOmniInputMode = "all", + inline_local_video: bool = False, + trust_remote_code: bool = False, + max_duration_seconds: float | None = None, + dataset_subset: str | None = None, + no_stream: bool = False, + **kwargs, + ) -> None: + if input_mode not in ("all", "visual", "audio"): + raise ValueError(f"input_mode must be 'all', 'visual', or 'audio', got {input_mode!r}") + + # Validate arguments: need either local JSON or HF path + if qa_json_path is None and dataset_path is None: + raise ValueError( + "Either 'qa_json_path' (local JSON) or 'dataset_path' (HuggingFace) must be provided. " + "For offline/air-gapped environments, download qa.json and use qa_json_path." + ) + + # Store configuration + self.qa_json_path = Path(qa_json_path) if qa_json_path else None + self.dataset_path = dataset_path + self.dataset_split = dataset_split + self.dataset_subset = dataset_subset + #: Match vLLM ``HuggingFaceDataset`` / bench CLI ``--no-stream``. + self._hf_streaming = not no_stream + self.video_dir = Path(video_dir) if video_dir else None + self.inline_local_video = inline_local_video + self.input_mode: DailyOmniInputMode = input_mode + self.max_duration_seconds = max_duration_seconds + self.trust_remote_code = trust_remote_code + + #: In-process cache of ffprobe durations only (no disk persistence). + self._video_durations: dict[str, float] = {} + + # Initialize parent BenchmarkDataset + super().__init__( + dataset_path=dataset_path if qa_json_path is None else None, + random_seed=random_seed, + **kwargs, + ) + + # Load data based on mode + self.load_data() + + # Verify dataset info + logger.info( + "Loaded Daily-Omni dataset: mode=%s, source=%s, random_seed=%d, input_mode=%s, max_duration=%s", + "local_json" if self.qa_json_path else "huggingface", + str(self.qa_json_path) if self.qa_json_path else f"{dataset_path}/{dataset_split}", + random_seed, + input_mode, + f"{max_duration_seconds}s" if max_duration_seconds else "unlimited", + ) + + def load_data(self) -> None: + """Populate ``self.data`` from either local JSON or the Hub. + + See module docstring: we do not subclass ``HuggingFaceDataset`` because Daily-Omni needs + a first-class offline path; Hub loading is an optional branch implemented below. + """ + if self.qa_json_path is not None: + self._load_from_local_json() + else: + self._load_from_huggingface() + + def _load_from_local_json(self) -> None: + """Load QA data from local JSON file.""" + if not self.qa_json_path.exists(): + raise FileNotFoundError(f"QA JSON file not found: {self.qa_json_path}") + + with open(self.qa_json_path, encoding="utf-8") as f: + data = json.load(f) + + # Support both list format and dict with "train"/"test" splits + if isinstance(data, dict): + # Try to get the requested split, fallback to first available + split_data = data.get(self.dataset_split) + if split_data is None: + available = list(data.keys()) + if available: + logger.warning( + "Split '%s' not found in %s, using '%s' instead", + self.dataset_split, + self.qa_json_path, + available[0], + ) + split_data = data[available[0]] + else: + split_data = [] + data = split_data + + if not isinstance(data, list): + raise ValueError(f"Expected list of QA items in JSON, got {type(data).__name__}") + + # Shuffle if requested + if not getattr(self, "disable_shuffle", False) and self.random_seed is not None: + import random + + rng = random.Random(self.random_seed) + shuffled = data[:] + rng.shuffle(shuffled) + data = shuffled + + # Create an iterator-like wrapper for compatibility + self.data = _ListDatasetIterator(data) + + def _load_from_huggingface(self) -> None: + """Load QA rows via ``datasets.load_dataset`` (legacy / convenience path). + + Kept for backward compatibility: callers can still pass ``dataset_path=liarliar/Daily-Omni`` + and get the same parquet-backed rows as the Hub dataset card, with streaming (or + non-streaming if ``no_stream=True``) and shuffle. + + This is intentionally **not** implemented by subclassing ``HuggingFaceDataset``: that base + always runs Hub ``load_dataset`` from its constructor and expects a Hub id as the primary + API; Daily-Omni instead chooses the source in ``load_data()`` (JSON vs Hub) while sharing + one ``sample()`` / request-building implementation for both. + """ + if load_dataset is None: + raise ImportError( + "datasets library is required for HuggingFace mode. " + "Install with: pip install datasets, or use local JSON mode instead." + ) + + ds = load_dataset( + self.dataset_path, + name=self.dataset_subset, + split=self.dataset_split, + streaming=self._hf_streaming, + trust_remote_code=self.trust_remote_code, + ) + if not getattr(self, "disable_shuffle", False): + ds = ds.shuffle(seed=self.random_seed) + self.data = ds + + def get_task_statistics(self) -> dict[str, int]: + """Get distribution of task types in the dataset. + + Returns: + Dict mapping task type to count + """ + stats: dict[str, int] = {} + for item in self.data: + row = self._coerce_row(item) + fields = self._normalize_qa_fields(row) + task_type = fields["task_type"] or "unknown" + stats[task_type] = stats.get(task_type, 0) + 1 + return stats + + @staticmethod + def _coerce_row(item: Any) -> dict[str, Any]: + """Turn a dataset row into a plain dict (Arrow / Mapping).""" + if isinstance(item, dict): + return item + if hasattr(item, "as_py"): + return dict(item.as_py()) # pyarrow Row + try: + return dict(item) + except (TypeError, ValueError): + return {k: item[k] for k in item} # type: ignore[misc] + + @staticmethod + def _normalize_qa_fields(row: dict[str, Any]) -> dict[str, Any]: + """Map official Daily-Omni qa.json / Hub schema to internal fields. + + Official fields (see liarliar/Daily-Omni ``qa.json``): ``Question``, ``Choice`` (list), + ``Answer``, ``video_id``, ``Type``, ``video_duration`` (``30s`` / ``60s``), ``video_category``, + plus other category columns. Legacy aliases (lowercase / older loaders) are still accepted. + """ + out: dict[str, Any] = {} + + out["question"] = str(row.get("Question") or row.get("question") or "").strip() + vid = row.get("video_id") if row.get("video_id") is not None else row.get("video") + out["video_id"] = str(vid).strip() if vid is not None else "" + out["task_type"] = str(row.get("Type") or row.get("task_type") or row.get("type") or "").strip() + vc = row.get("video_category") if row.get("video_category") is not None else row.get("videoCategory") + out["video_category"] = str(vc).strip() if vc is not None else "" + vd = row.get("video_duration") if row.get("video_duration") is not None else row.get("videoDuration") + out["video_duration"] = str(vd).strip() if vd is not None else "" + out["answer"] = str(row.get("Answer") or row.get("answer") or "").strip() + vu = row.get("video_url") if row.get("video_url") is not None else row.get("Video_URL") + out["video_url"] = str(vu).strip() if vu is not None and str(vu).strip() else None + + choice = row.get("Choice") + if choice is None: + choice = row.get("options") or row.get("choice") + out["choice"] = choice + + return out + + def sample( + self, + tokenizer: TokenizerLike, + num_requests: int, + output_len: int | None = None, + request_id_prefix: str = "", + no_oversample: bool = False, + **kwargs, + ) -> list[SampleRequest]: + """Sample requests from Daily-Omni dataset. + + Args: + tokenizer: Tokenizer for computing prompt length + num_requests: Number of requests to sample + output_len: Target output length in tokens (default: 256) + request_id_prefix: Prefix for request IDs + no_oversample: If True, do not oversample if fewer examples available + **kwargs: Additional arguments (ignored) + + Returns: + List of SampleRequest objects with video URLs and prompts + """ + if output_len is None: + output_len = self.DEFAULT_OUTPUT_LEN + + sampled_requests: list[SampleRequest] = [] + ind = 0 + cached_tokenizer = get_cached_tokenizer(tokenizer) + + # Iterate over shuffled dataset + for item in self.data: + if len(sampled_requests) >= num_requests: + break + + request = self._create_sample_request( + self._coerce_row(item), cached_tokenizer, output_len, request_id_prefix, ind + ) + if request: + sampled_requests.append(request) + ind += 1 + + logger.info("Created %d sample requests from Daily-Omni dataset", len(sampled_requests)) + + # Handle oversampling if needed + self.maybe_oversample_requests(sampled_requests, num_requests, request_id_prefix, no_oversample) + + return sampled_requests + + def _create_sample_request( + self, + qa_item: dict[str, Any], + tokenizer: TokenizerLike, + output_len: int, + request_id_prefix: str, + index: int, + ) -> SampleRequest | None: + """Create a SampleRequest from a QA item. + + Args: + qa_item: QA pair from the dataset + tokenizer: Tokenizer + output_len: Target output length + request_id_prefix: Prefix for request ID + index: Request index + + Returns: + SampleRequest or None if invalid + """ + fields = self._normalize_qa_fields(qa_item) + video_id = fields["video_id"] + question = fields["question"] + choice = fields["choice"] + task_type = fields["task_type"] + video_url = fields["video_url"] + video_duration = fields.get("video_duration") or "" + video_category = fields.get("video_category") or "" + + if not video_id and not video_url: + logger.warning("Skipping item: no video_id / video_url") + return None + + if not question: + logger.warning("Skipping item: no question found") + return None + + # Official layout after extracting Videos.tar (see Lliar-liar/Daily-Omni test_model): + # {video_base_dir}/{video_id}/{video_id}_video.mp4 + mm_payload, omni_extra, mm_pos = self._compose_daily_omni_multimodal(video_id, video_url) + if not mm_payload: + return None + + messages = self._build_daily_omni_openai_messages(mm_payload, question, choice) + user_text = self._official_daily_omni_user_prompt(question, choice) + # Text-only length estimate (same as before: no MM token count in bench). + prompt_len = len(tokenizer.encode(f"{DAILY_OMNI_SYSTEM_TEXT}\n{user_text}")) + + return DailyOmniSampleRequest( + prompt=user_text, + prompt_len=prompt_len, + expected_output_len=output_len, + multi_modal_data=None, + request_id=f"{request_id_prefix}{index}", + daily_omni_gold_answer=fields["answer"], + daily_omni_video_id=video_id, + daily_omni_task_type=task_type, + daily_omni_video_duration=video_duration, + daily_omni_video_category=video_category, + omni_extra_body=omni_extra, + omni_chat_messages=messages, + omni_chat_mm_position=mm_pos, + ) + + @staticmethod + def _official_video_relpath(video_id: str) -> str: + """Relative path inside extracted ``Videos/`` per upstream Daily-Omni scripts.""" + return f"{video_id}/{video_id}_video.mp4" + + @staticmethod + def _official_audio_relpath(video_id: str) -> str: + """Relative path for extracted WAV per upstream ``get_audio_path``.""" + return f"{video_id}/{video_id}_audio.wav" + + def _resolve_local_video_path(self, video_id: str) -> Path | None: + """Pick an existing file under ``video_dir`` (official layout + flat fallback).""" + if not self.video_dir or not video_id: + return None + + candidates = [ + self.video_dir / self._official_video_relpath(video_id), + self.video_dir / f"{video_id}.mp4", # flat layout (custom mirrors / outdated docs) + ] + seen: set[Path] = set() + for p in candidates: + rp = p.resolve() + if rp in seen: + continue + seen.add(rp) + if p.exists(): + return p + return None + + def _resolve_local_audio_path(self, video_id: str) -> Path | None: + """Pick an existing WAV under ``video_dir`` (official layout + flat fallback).""" + if not self.video_dir or not video_id: + return None + candidates = [ + self.video_dir / self._official_audio_relpath(video_id), + self.video_dir / f"{video_id}.wav", + ] + seen: set[Path] = set() + for p in candidates: + rp = p.resolve() + if rp in seen: + continue + seen.add(rp) + if p.exists(): + return p + return None + + def _local_file_to_video_url_payload(self, video_path: Path) -> dict[str, Any]: + """Build OpenAI-style video_url part for a resolved local file. + + vLLM rejects ``file://`` unless the server was started with + ``--allowed-local-media-path`` set to a directory that **contains** the file + (typically the extracted ``Videos`` root). Use ``inline_local_video=True`` to + send base64 data URLs instead (no server path allowlist; larger requests). + """ + path = video_path.expanduser().resolve() + if self.inline_local_video: + raw = path.read_bytes() + b64 = base64.b64encode(raw).decode("ascii") + return { + "type": "video_url", + "video_url": {"url": f"data:video/mp4;base64,{b64}"}, + } + return { + "type": "video_url", + "video_url": {"url": path.as_uri()}, + } + + def _local_file_to_audio_url_payload(self, audio_path: Path) -> dict[str, Any]: + """Build OpenAI-style ``audio_url`` part for a resolved local WAV file.""" + path = audio_path.expanduser().resolve() + if self.inline_local_video: + raw = path.read_bytes() + b64 = base64.b64encode(raw).decode("ascii") + return { + "type": "audio_url", + "audio_url": {"url": f"data:audio/wav;base64,{b64}"}, + } + return { + "type": "audio_url", + "audio_url": {"url": path.as_uri()}, + } + + def _get_video_content( + self, + video_id: str, + video_url: str | None, + ) -> dict[str, Any] | None: + """Resolve video for OpenAI-style ``video_url`` content. + + Upstream uses ``get_video_path(video_id, base) -> base/video_id/video_id_video.mp4``. + The Hub repo only publishes ``Videos.tar``; use ``--daily-omni-video-dir`` pointing + at the extracted ``Videos`` folder (parent of per-``video_id`` subdirs). + + For ``file://`` URLs, start ``vllm serve`` with e.g. + ``--allowed-local-media-path /same/path/as/daily-omni-video-dir``. + """ + if video_url: + url = video_url + if not url.startswith(("http://", "https://", "file://")): + url = f"https://{url.lstrip('/')}" + return {"type": "video_url", "video_url": {"url": url}} + + if self.video_dir and video_id: + video_path = self._resolve_local_video_path(video_id) + if video_path is not None: + return self._local_file_to_video_url_payload(video_path) + logger.warning( + "Video not found under video_dir=%s for video_id=%r (expected %s or %s)", + self.video_dir, + video_id, + self._official_video_relpath(video_id), + f"{video_id}.mp4", + ) + + if video_id: + repo = self.dataset_path or self.DEFAULT_HF_DATASET_ID + rel = self._official_video_relpath(video_id) + hf_video_url = f"https://huggingface.co/datasets/{repo}/resolve/main/Videos/{rel}" + logger.debug( + "Using HF video URL (likely 404 — Hub ships Videos.tar only): %s", + hf_video_url, + ) + return {"type": "video_url", "video_url": {"url": hf_video_url}} + + logger.error("Could not determine video source for video_id=%r", video_id) + return None + + def _get_audio_content(self, video_id: str) -> dict[str, Any] | None: + """Resolve extracted WAV for OpenAI-style ``audio_url`` (local files only).""" + if not self.video_dir or not video_id: + logger.warning( + "Daily-Omni input_mode %r requires --daily-omni-video-dir with %s", + self.input_mode, + self._official_audio_relpath(video_id), + ) + return None + audio_path = self._resolve_local_audio_path(video_id) + if audio_path is not None: + return self._local_file_to_audio_url_payload(audio_path) + logger.warning( + "Audio not found under video_dir=%s for video_id=%r (expected %s or %s)", + self.video_dir, + video_id, + self._official_audio_relpath(video_id), + f"{video_id}.wav", + ) + return None + + def _compose_daily_omni_multimodal( + self, + video_id: str, + video_url: str | None, + ) -> tuple[dict[str, Any] | list[dict[str, Any]] | None, dict[str, Any] | None, Literal["first", "last"]]: + """Build ``multi_modal_data`` and request extras for the active ``input_mode``. + + Mirrors upstream Daily-Omni: separate video + WAV with ``use_audio_in_video=False``. + """ + extra: dict[str, Any] = {"mm_processor_kwargs": {"use_audio_in_video": False}} + mode = self.input_mode + + if mode == "visual": + v = self._get_video_content(video_id, video_url) + return v, extra, "last" + + if mode == "audio": + a = self._get_audio_content(video_id) + return a, extra, "first" + + v = self._get_video_content(video_id, video_url) + a = self._get_audio_content(video_id) + if not v or not a: + return None, None, "first" + return [v, a], extra, "first" + + @staticmethod + def _media_desc_for_official_prompt(mode: DailyOmniInputMode) -> str: + """``media_desc`` in upstream ``build_conversation``.""" + if mode == "audio": + return "given audio" + if mode == "all": + return "given video and audio together" + return "given video" + + @staticmethod + def _choices_repr_for_official_prompt(choice: Any) -> str: + """Format ``Choice`` from qa.json for the model (one option per line when possible). + + Using ``str(list)`` embeds Python list brackets and quotes, which is poor for MCQ + reading; lists/tuples are joined with newlines instead. Other shapes fall back to + ``str(choice)`` for parity with exotic upstream payloads. + """ + if choice is None: + return "" + if isinstance(choice, (list, tuple)): + lines = [str(x).strip() for x in choice if str(x).strip()] + return "\n".join(lines) + if isinstance(choice, dict): + return "\n".join(f"{k}. {v}" for k, v in choice.items()) + return str(choice) + + def _official_daily_omni_user_prompt(self, question: str, choice: Any) -> str: + """User text block from Daily-Omni ``build_conversation`` (after media parts).""" + task_prompt = self._media_desc_for_official_prompt(self.input_mode) + choices = self._choices_repr_for_official_prompt(choice) + # Single f-string with explicit newlines avoids accidental implicit concatenation + # gluing sentences (e.g. ``...media_desc.Select...``) when editing. + return ( + "Your task is to accurately answer multiple-choice questions " + f"based on the {task_prompt}.\n" + "Select the single most accurate answer from the given choices.\n" + f"Question: {question}\n" + f"Choices: {choices}\n" + "Your answer should be a capital letter representing your choice: " + "A, B, C, or D. Don't generate any other text.\n" + ) + + def _build_daily_omni_openai_messages( + self, + mm_payload: dict[str, Any] | list[dict[str, Any]], + question: str, + choice: Any, + ) -> list[dict[str, Any]]: + """Map upstream conversation to OpenAI Chat Completions ``messages`` (video_url / audio_url parts).""" + user_text = self._official_daily_omni_user_prompt(question, choice) + mm_list: list[dict[str, Any]] = mm_payload if isinstance(mm_payload, list) else [mm_payload] + user_content: list[dict[str, Any]] = [*mm_list, {"type": "text", "text": user_text}] + return [ + {"role": "system", "content": [{"type": "text", "text": DAILY_OMNI_SYSTEM_TEXT}]}, + {"role": "user", "content": user_content}, + ] + + def sample_by_task_type( + self, + tokenizer: TokenizerLike, + task_type: str, + num_samples: int, + output_len: int | None = None, + request_id_prefix: str = "", + **kwargs, + ) -> list[SampleRequest]: + """Sample requests filtered by task type. + + Args: + tokenizer: Tokenizer + task_type: Task type to filter by + num_samples: Number of samples + output_len: Target output length + request_id_prefix: Prefix for request IDs + **kwargs: Additional sampling arguments + + Returns: + List of SampleRequest objects matching the task type + """ + if output_len is None: + output_len = self.DEFAULT_OUTPUT_LEN + + filtered = [ + item for item in self.data if self._normalize_qa_fields(self._coerce_row(item))["task_type"] == task_type + ] + + available = len(filtered) + if available < num_samples: + logger.warning( + "Only %d samples available for task type '%s', requested %d", + available, + task_type, + num_samples, + ) + num_samples = available + + sampled_requests: list[SampleRequest] = [] + cached_tokenizer = get_cached_tokenizer(tokenizer) + + for i, item in enumerate(filtered[:num_samples]): + request = self._create_sample_request(item, cached_tokenizer, output_len, request_id_prefix, i) + if request: + sampled_requests.append(request) + + return sampled_requests + + def __repr__(self) -> str: + return ( + f"DailyOmniDataset(" + f"dataset_path={self.dataset_path!r}, " + f"dataset_split={self.dataset_split!r}, " + f"video_dir={self.video_dir!r}, " + f"input_mode={self.input_mode!r}, " + f"inline_local_video={self.inline_local_video!r}, " + f"max_duration_seconds={self.max_duration_seconds}, " + f"random_seed={self.random_seed}" + f")" + ) + + +def load_daily_omni_dataset( + qa_json_path: str | None = None, + dataset_path: str | None = None, + dataset_split: str = "train", + random_seed: int = 0, + video_dir: str | None = None, + input_mode: DailyOmniInputMode = "all", + max_duration_seconds: float | None = None, + dataset_subset: str | None = None, + no_stream: bool = False, + **kwargs, +) -> DailyOmniDataset: + """Convenience function to load Daily-Omni dataset. + + Args: + qa_json_path: Path to local qa.json file (recommended for offline/air-gapped environments). + When provided, ``dataset_path`` is ignored. + dataset_path: HuggingFace dataset path (default: liarliar/Daily-Omni). Used only if + ``qa_json_path`` is not provided (legacy online mode). + dataset_split: Dataset split to use (default: "train") + random_seed: Random seed for shuffling + video_dir: Directory containing extracted ``Videos/`` tree (MP4 and, for ``all``/``audio``, WAV) + input_mode: ``visual`` | ``audio`` | ``all`` (same semantics as upstream Daily-Omni) + max_duration_seconds: Maximum video duration in seconds (e.g., 30 for 30s subset, 60 for 60s subset); + uses ffprobe on local files under ``video_dir`` (in-memory cache only for this process). + **kwargs: Additional arguments passed to DailyOmniDataset + + Returns: + DailyOmniDataset instance + + Example: + >>> from vllm_omni.benchmarks.data_modules.daily_omni_dataset import load_daily_omni_dataset + + # Local JSON mode (recommended for offline) + >>> dataset = load_daily_omni_dataset( + ... qa_json_path="/path/to/qa.json", + ... video_dir="/path/to/Daily-Omni/Videos", + ... random_seed=42, + ... max_duration_seconds=30, + ... ) + + # HuggingFace mode (legacy online) + >>> dataset = load_daily_omni_dataset( + ... dataset_path="liarliar/Daily-Omni", + ... video_dir="/path/to/Daily-Omni/Videos", + ... random_seed=42, + ... ) + >>> requests = dataset.sample(tokenizer, num_requests=100) + """ + return DailyOmniDataset( + qa_json_path=qa_json_path, + dataset_path=dataset_path, + dataset_split=dataset_split, + random_seed=random_seed, + video_dir=video_dir, + input_mode=input_mode, + max_duration_seconds=max_duration_seconds, + dataset_subset=dataset_subset, + no_stream=no_stream, + **kwargs, + ) + + +def get_daily_omni_statistics( + qa_json_path: str | None = None, + dataset_path: str | None = DailyOmniDataset.DEFAULT_HF_DATASET_ID, + dataset_split: str = "train", +) -> dict[str, Any]: + """Get statistics about the Daily-Omni dataset. + + Args: + qa_json_path: Path to local qa.json file (recommended for offline/air-gapped environments). + When provided, ``dataset_path`` is ignored. + dataset_path: HuggingFace dataset path. Defaults to ``DailyOmniDataset.DEFAULT_HF_DATASET_ID`` + when ``qa_json_path`` is omitted. Pass ``None`` only together with ``qa_json_path``. + dataset_split: Dataset split to use (default: "train") + + Returns: + Statistics dict with task type distribution and other info + + Example: + >>> from vllm_omni.benchmarks.data_modules.daily_omni_dataset import get_daily_omni_statistics + + # Local JSON mode + >>> stats = get_daily_omni_statistics(qa_json_path="/path/to/qa.json") + + # HuggingFace mode + >>> stats = get_daily_omni_statistics(dataset_path="liarliar/Daily-Omni") + >>> print(f"Total QA pairs: {stats['total_qa_pairs']}") + >>> print(f"Task distribution: {stats['task_distribution']}") + """ + dataset = DailyOmniDataset( + qa_json_path=qa_json_path, + dataset_path=dataset_path, + dataset_split=dataset_split, + ) + task_stats = dataset.get_task_statistics() + + source = str(qa_json_path) if qa_json_path else f"{dataset_path}/{dataset_split}" + return { + "source": source, + "total_qa_pairs": len(list(dataset.data)), + "task_distribution": task_stats, + } diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_eval.py b/vllm_omni/benchmarks/data_modules/daily_omni_eval.py new file mode 100644 index 000000000..ecc9edc84 --- /dev/null +++ b/vllm_omni/benchmarks/data_modules/daily_omni_eval.py @@ -0,0 +1,406 @@ +"""Daily-Omni multiple-choice accuracy scoring for vLLM-Omni bench serve. + +Compares model ``generated_text`` to dataset ``Answer`` (A/B/C/D). + +**Alignment with open-source** (`Lliar-liar/Daily-Omni` ``test_model/.../testmodel.py``): + +- Answer extraction defaults to the same rules as ``extract_choice_letter`` (strip after an + ``assistant`` marker, then leading ``A``–``D``, else first ``\\b[A-D]\\b``). Set env + ``DAILY_OMNI_EXTRACT_MODE=relaxed`` to use the older vLLM-Omni heuristics (last ``answer:``, + tail scan, etc.). +- Overall accuracy comparable to the official script uses **successful HTTP responses only** as + the denominator (their ``valid_questions = total - failed`` excludes inference / I/O skips). + We also report ``daily_omni_accuracy_incl_http_fail`` where each failed request counts as a + wrong answer in the denominator (stricter throughput-bench view). +- **By video length:** mirrors upstream ``--- Accuracy by Video Duration ---`` for ``30s`` / + ``60s`` (``qa.json`` ``video_duration``): ``daily_omni_per_duration*`` metrics and a printed block. +- **By video category:** mirrors ``--- Accuracy by Video Category ---`` using ``video_category`` + from ``qa.json`` (``daily_omni_per_category*``; empty category is bucketed as ``unknown``). +- **Correctness:** uses the same ``evaluate_answer`` rule as upstream (truthy extracted letter vs + raw ``Answer`` string, both ``strip().upper()``). Rows with empty ``Answer`` are skipped + (``no_gold``), matching missing-field skips in the official loop. +""" + +from __future__ import annotations + +import os +import re +from typing import Any + +from vllm.benchmarks.lib.endpoint_request_func import RequestFuncOutput + +from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniSampleRequest + +_VALID = frozenset("ABCD") + +# Official ``testmodel.py`` buckets (``qa.json`` ``video_duration``). +DAILY_OMNI_DURATION_KEYS: tuple[str, ...] = ("30s", "60s") + + +def extract_choice_letter_official(text: str | None) -> str | None: + """Port of Daily-Omni ``extract_choice_letter`` (first A–D, assistant-tail semantics).""" + if not text: + return None + raw = str(text).strip() + if not raw: + return None + match = re.search(r"assistant\s*([\s\S]*)$", raw, flags=re.IGNORECASE) + candidate = match.group(1).strip() if match else raw + direct = re.match(r"(?i)^\s*([A-D])(?:[\s\.\):：]|$)", candidate) + if direct: + return direct.group(1).upper() + fallback = re.search(r"\b([A-D])\b", candidate.upper()) + if fallback: + return fallback.group(1) + return None + + +def evaluate_answer_official(model_answer: str | None, correct_answer: str) -> bool: + """Port of Daily-Omni ``evaluate_answer`` (strict string match after strip/upper).""" + if not model_answer: + return False + return model_answer.strip().upper() == (correct_answer or "").strip().upper() + + +def normalize_gold_answer(gold: str) -> str | None: + """Best-effort single letter from ``Answer`` (for ``gold_normalized`` in saved items only).""" + g = (gold or "").strip().upper() + if len(g) == 1 and g in _VALID: + return g + m = re.search(r"([ABCD])\b", g) + if m: + return m.group(1).upper() + return None + + +def _extract_predicted_choice_relaxed(text: str) -> str | None: + """Legacy vLLM-Omni heuristics (last ``answer:`` patterns, tail scan).""" + if not text or not str(text).strip(): + return None + t = str(text).strip() + + strong_patterns = [ + r"(?i)\*\*answer\*\*\s*[:：]?\s*\(?([ABCD])\)?", + r"(?i)\banswer\s*[:：]?\s*\(?([ABCD])\)?", + r"(?i)\bfinal\s+answer\s*[:：]?\s*\(?([ABCD])\)?", + r"(?i)\bcorrect\s+(?:answer|option)\s*[:：]?\s*\(?([ABCD])\)?", + r"(?i)\bthe\s+(?:correct\s+)?option\s+(?:is|would\s+be)\s*\(?([ABCD])\)?", + r"(?i)\bI\s+(?:would\s+)?(?:choose|select|pick)\s*\(?([ABCD])\)?", + ] + last_letter: str | None = None + for pat in strong_patterns: + for m in re.finditer(pat, t): + last_letter = m.group(1).upper() + if last_letter: + return last_letter + + # Weaker phrases: first match can be spurious; still prefer last occurrence. + weak_patterns = [ + r"(?i)\boption\s*[:：]?\s*\(?([ABCD])\)?", + r"(?i)\bchoice\s*[:：]?\s*\(?([ABCD])\)?", + ] + for pat in weak_patterns: + for m in re.finditer(pat, t): + last_letter = m.group(1).upper() + if last_letter: + return last_letter + + paren = list(re.finditer(r"\(([ABCD])\)", t)) + if paren: + return paren[-1].group(1).upper() + + # First line sometimes is just "B" or "B." — allow if whole output is short + one_line = t.split("\n", 1)[0].strip() + if len(t) < 120 and len(one_line) <= 6: + m0 = re.match(r"^([ABCD])\s*[.:\)]?\s*$", one_line, re.I) + if m0: + return m0.group(1).upper() + + # Tail-only: avoids matching echoed "A. ..." option blocks at the start + tail_len = min(500, len(t)) + tail = t[-tail_len:] + # ``\b`` after the letter avoids "Because"/"Definitely" false positives + m = re.search(r"(?:^|[^\w])([ABCD])\b", tail, re.I) + if m: + return m.group(1).upper() + + return None + + +def extract_predicted_choice(text: str | None) -> str | None: + """Parse model output to A–D (official Daily-Omni rules by default).""" + if not text or not str(text).strip(): + return None + mode = os.environ.get("DAILY_OMNI_EXTRACT_MODE", "official").strip().lower() + if mode in ("relaxed", "heuristic", "legacy"): + return _extract_predicted_choice_relaxed(str(text)) + return extract_choice_letter_official(text) + + +def compute_daily_omni_accuracy_metrics( + input_requests: list[Any], + outputs: list[RequestFuncOutput], + *, + include_per_item: bool = False, +) -> dict[str, Any] | None: + """If all requests are :class:`DailyOmniSampleRequest`, compute accuracy stats. + + Rows with empty ``Answer`` (after strip) are skipped as ``no_gold``, like upstream missing + ``correct_answer``. + + **Denominators:** The open-source script excludes items that hit inference / I/O failures + from ``valid_questions``; we mirror that with ``daily_omni_accuracy`` (= correct / + successful responses). Failed HTTP requests are also tracked and used in + ``daily_omni_accuracy_incl_http_fail`` (each failure counts as incorrect in the + denominator). + """ + if not input_requests or len(input_requests) != len(outputs): + return None + if not all(isinstance(r, DailyOmniSampleRequest) for r in input_requests): + return None + + # total / correct: all rows with gold (incl. HTTP fail in total) + # total_ok / correct_ok: successful HTTP only (GitHub-style per-type denominator) + per_task: dict[str, dict[str, int]] = {} + per_category: dict[str, dict[str, int]] = {} + per_duration: dict[str, dict[str, int]] = { + k: {"correct": 0, "total": 0, "correct_ok": 0, "total_ok": 0} for k in DAILY_OMNI_DURATION_KEYS + } + items: list[dict[str, Any]] = [] + correct = 0 + evaluated = 0 + no_gold = 0 + request_failed = 0 + parse_failed = 0 # success but could not extract A–D + + for req, out in zip(input_requests, outputs, strict=True): + assert isinstance(req, DailyOmniSampleRequest) + gold_raw = (req.daily_omni_gold_answer or "").strip() + gold_norm = normalize_gold_answer(req.daily_omni_gold_answer) + tt = (req.daily_omni_task_type or "unknown").strip() or "unknown" + dur_key = (req.daily_omni_video_duration or "").strip() + dur_active = dur_key in per_duration + cat_key = (req.daily_omni_video_category or "").strip() or "unknown" + if tt not in per_task: + per_task[tt] = {"correct": 0, "total": 0, "correct_ok": 0, "total_ok": 0} + if cat_key not in per_category: + per_category[cat_key] = {"correct": 0, "total": 0, "correct_ok": 0, "total_ok": 0} + + if not gold_raw: + no_gold += 1 + items.append( + { + "request_id": req.request_id, + "skipped": True, + "reason": "no_gold", + "task_type": tt, + "video_id": req.daily_omni_video_id, + "video_duration": dur_key or None, + "video_category": cat_key if cat_key != "unknown" else None, + } + ) + continue + + if not out.success: + request_failed += 1 + evaluated += 1 + per_task[tt]["total"] += 1 + per_category[cat_key]["total"] += 1 + if dur_active: + per_duration[dur_key]["total"] += 1 + # GitHub: failed inference not in valid_questions — do not increment total_ok + items.append( + { + "request_id": req.request_id, + "gold": gold_raw, + "gold_normalized": gold_norm, + "predicted": None, + "correct": False, + "task_type": tt, + "video_id": req.daily_omni_video_id, + "video_duration": dur_key or None, + "video_category": cat_key if cat_key != "unknown" else None, + "error": (out.error or "")[:500], + } + ) + continue + + pred = extract_predicted_choice(out.generated_text) + evaluated += 1 + per_task[tt]["total"] += 1 + per_task[tt]["total_ok"] += 1 + per_category[cat_key]["total"] += 1 + per_category[cat_key]["total_ok"] += 1 + if dur_active: + per_duration[dur_key]["total"] += 1 + per_duration[dur_key]["total_ok"] += 1 + if pred is None: + parse_failed += 1 + is_correct = evaluate_answer_official(pred, req.daily_omni_gold_answer) + if is_correct: + correct += 1 + per_task[tt]["correct"] += 1 + per_task[tt]["correct_ok"] += 1 + per_category[cat_key]["correct"] += 1 + per_category[cat_key]["correct_ok"] += 1 + if dur_active: + per_duration[dur_key]["correct"] += 1 + per_duration[dur_key]["correct_ok"] += 1 + + items.append( + { + "request_id": req.request_id, + "gold": gold_raw, + "gold_normalized": gold_norm, + "predicted": pred, + "correct": is_correct, + "parse_failed": pred is None, + "task_type": tt, + "video_id": req.daily_omni_video_id, + "video_duration": dur_key or None, + "video_category": cat_key if cat_key != "unknown" else None, + } + ) + + evaluated_ok = evaluated - request_failed + accuracy_github = (correct / evaluated_ok) if evaluated_ok else None + accuracy_incl_fail = (correct / evaluated) if evaluated else None + + per_task_accuracy: dict[str, float | None] = {} + per_task_accuracy_github: dict[str, float | None] = {} + for name, st in per_task.items(): + tot = st["total"] + per_task_accuracy[name] = (st["correct"] / tot) if tot else None + tok = st["total_ok"] + per_task_accuracy_github[name] = (st["correct_ok"] / tok) if tok else None + + per_category_accuracy: dict[str, float | None] = {} + per_category_accuracy_github: dict[str, float | None] = {} + for name, st in per_category.items(): + tot = st["total"] + per_category_accuracy[name] = (st["correct"] / tot) if tot else None + tok = st["total_ok"] + per_category_accuracy_github[name] = (st["correct_ok"] / tok) if tok else None + + per_duration_accuracy: dict[str, float | None] = {} + per_duration_accuracy_github: dict[str, float | None] = {} + for name, st in per_duration.items(): + tot = st["total"] + per_duration_accuracy[name] = (st["correct"] / tot) if tot else None + tok = st["total_ok"] + per_duration_accuracy_github[name] = (st["correct_ok"] / tok) if tok else None + + out: dict[str, Any] = { + # Comparable to GitHub testmodel.py: correct / successful inferences + "daily_omni_accuracy": accuracy_github, + "daily_omni_accuracy_incl_http_fail": accuracy_incl_fail, + "daily_omni_correct": correct, + "daily_omni_evaluated": evaluated, + "daily_omni_evaluated_ok": evaluated_ok, + "daily_omni_no_gold": no_gold, + "daily_omni_request_failed": request_failed, + "daily_omni_parse_failed": parse_failed, + "daily_omni_per_task": {k: dict(v) for k, v in per_task.items()}, + "daily_omni_per_task_accuracy": per_task_accuracy, + "daily_omni_per_task_accuracy_github_style": per_task_accuracy_github, + "daily_omni_per_category": {k: dict(v) for k, v in per_category.items()}, + "daily_omni_per_category_accuracy": per_category_accuracy, + "daily_omni_per_category_accuracy_github_style": per_category_accuracy_github, + "daily_omni_per_duration": {k: dict(v) for k, v in per_duration.items()}, + "daily_omni_per_duration_accuracy": per_duration_accuracy, + "daily_omni_per_duration_accuracy_github_style": per_duration_accuracy_github, + } + if include_per_item: + out["daily_omni_eval_items"] = items + return out + + +def print_daily_omni_accuracy_summary(metrics: dict[str, Any]) -> None: + """Pretty-print accuracy block (stdout).""" + acc = metrics.get("daily_omni_accuracy") + acc_fail = metrics.get("daily_omni_accuracy_incl_http_fail") + if acc is None and acc_fail is None and metrics.get("daily_omni_evaluated", 0) == 0: + return + print("{s:{c}^{n}}".format(s=" Daily-Omni accuracy (MCQ) ", n=50, c="=")) + ok = int(metrics.get("daily_omni_evaluated_ok", 0) or 0) + cor = int(metrics.get("daily_omni_correct", 0) or 0) + if ok > 0 and acc is not None: + print(f"Overall Accuracy: {cor}/{ok} = {acc:.2%}") + elif int(metrics.get("daily_omni_evaluated", 0) or 0) > 0: + print("Overall Accuracy: 0/0 = N/A (no successful HTTP responses)") + print( + "{:<40} {:<10}".format( + "Submitted (gold present):", + metrics.get("daily_omni_evaluated", 0), + ) + ) + print( + "{:<40} {:<10}".format( + "Successful HTTP (GitHub denom.):", + metrics.get("daily_omni_evaluated_ok", 0), + ) + ) + print("{:<40} {:<10}".format("Correct:", metrics.get("daily_omni_correct", 0))) + if acc is not None: + print("{:<40} {:<10.4f}".format("Accuracy (ratio, same as above):", acc)) + if acc_fail is not None and metrics.get("daily_omni_request_failed", 0): + print( + "{:<40} {:<10.4f}".format( + "Accuracy (incl. HTTP as wrong):", + acc_fail, + ) + ) + print("{:<40} {:<10}".format("Skipped (no gold):", metrics.get("daily_omni_no_gold", 0))) + print( + "{:<40} {:<10}".format( + "HTTP failed (excl. from GitHub acc.):", + metrics.get("daily_omni_request_failed", 0), + ) + ) + print( + "{:<40} {:<10}".format( + "Parsed OK but no A–D found:", + metrics.get("daily_omni_parse_failed", 0), + ) + ) + pt = metrics.get("daily_omni_per_task") or {} + pta = metrics.get("daily_omni_per_task_accuracy_github_style") or {} + if pta: + print("\n--- Accuracy by QA Type ---") + for name in sorted(pta.keys()): + a = pta[name] + st = pt.get(name) or {} + tok = int(st.get("total_ok", 0) or 0) + cok = int(st.get("correct_ok", 0) or 0) + if tok and a is not None: + print(f"{name}: {cok}/{tok} = {a:.2%}") + else: + print(f"{name}: 0/0 = N/A") + + pc = metrics.get("daily_omni_per_category") or {} + ptc = metrics.get("daily_omni_per_category_accuracy_github_style") or {} + if ptc: + print("\n--- Accuracy by Video Category ---") + for name in sorted(ptc.keys()): + a = ptc[name] + st = pc.get(name) or {} + tok = int(st.get("total_ok", 0) or 0) + cok = int(st.get("correct_ok", 0) or 0) + if tok and a is not None: + print(f"{name}: {cok}/{tok} = {a:.2%}") + else: + print(f"{name}: 0/0 = N/A") + + pdf = metrics.get("daily_omni_per_duration_accuracy_github_style") or {} + if pdf: + print("\n--- Accuracy by Video Duration ---") + for name in DAILY_OMNI_DURATION_KEYS: + a = pdf.get(name) + st = (metrics.get("daily_omni_per_duration") or {}).get(name) or {} + tok = int(st.get("total_ok", 0) or 0) + cor = int(st.get("correct_ok", 0) or 0) + if tok and a is not None: + print(f"{name} Duration: {cor}/{tok} = {a:.2%}") + else: + print(f"{name} Duration: 0/0 = N/A") + print("=" * 50) diff --git a/vllm_omni/benchmarks/data_modules/daily_omni_text_audio.py b/vllm_omni/benchmarks/data_modules/daily_omni_text_audio.py new file mode 100644 index 000000000..69fbe026b --- /dev/null +++ b/vllm_omni/benchmarks/data_modules/daily_omni_text_audio.py @@ -0,0 +1,255 @@ +"""Daily-Omni: optional consistency check between text stream and generated speech. + +The benchmark MCQ accuracy uses ``generated_text`` only. When the omni server also +streams ``modality=audio`` (TTS), this module can transcribe the concatenated WAV +with Whisper and compare the inferred option letter to the one parsed from text. + +Requires ``openai-whisper`` (``pip install openai-whisper``). Enable via env +``DAILY_OMNI_TEXT_AUDIO_CONSISTENCY=1`` or CLI ``--daily-omni-text-audio-consistency``. + +Whisper model name defaults to ``tiny`` (override with ``DAILY_OMNI_WHISPER_MODEL``). +""" + +from __future__ import annotations + +import logging +import os +import re +import threading +from typing import Any + +from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniSampleRequest +from vllm_omni.benchmarks.data_modules.daily_omni_eval import extract_predicted_choice + +logger = logging.getLogger(__name__) + +_whisper_model = None +_whisper_model_name: str | None = None +_whisper_lock = threading.Lock() + + +def env_text_audio_check_enabled() -> bool: + return os.environ.get("DAILY_OMNI_TEXT_AUDIO_CONSISTENCY", "").lower() in ( + "1", + "true", + "yes", + ) + + +def extract_choice_from_asr_transcript(transcript: str) -> str | None: + """Parse A–D from ASR text; extends :func:`extract_predicted_choice` with spoken Chinese phrases.""" + c = extract_predicted_choice(transcript) + if c: + return c + t = transcript or "" + for pat in ( + r"(?i)选项\s*([ABCD])\b", + r"(?i)选\s*([ABCD])\b", + r"(?i)答案\s*是\s*([ABCD])\b", + r"(?i)答案\s*([ABCD])\b", + ): + m = re.search(pat, t) + if m: + return m.group(1).upper() + return None + + +def _get_whisper_model(model_name: str): + global _whisper_model, _whisper_model_name + with _whisper_lock: + if _whisper_model is None or _whisper_model_name != model_name: + import whisper + + logger.warning( + "Loading Whisper model %r for Daily-Omni text/audio consistency (one-time)...", + model_name, + ) + _whisper_model = whisper.load_model(model_name) + _whisper_model_name = model_name + return _whisper_model + + +def transcribe_wav_bytes( + wav_bytes: bytes, + *, + language: str | None = None, + model_name: str | None = None, +) -> tuple[str | None, str | None]: + """Transcribe WAV bytes. Returns ``(transcript, error)`` — one of them is set. + + Args: + wav_bytes: RIFF WAV file bytes. + language: Optional Whisper language code (e.g. ``en``, ``zh``); improves accuracy/latency. + model_name: Override model id; else ``DAILY_OMNI_WHISPER_MODEL`` or ``tiny``. + """ + if not wav_bytes: + return None, "empty_wav" + if model_name is None or not str(model_name).strip(): + model_name = os.environ.get("DAILY_OMNI_WHISPER_MODEL") or "tiny" + model_name = str(model_name).strip() or "tiny" + path: str | None = None + try: + import tempfile + + model = _get_whisper_model(model_name) + fd, path = tempfile.mkstemp(suffix=".wav") + with os.fdopen(fd, "wb") as fp: + fp.write(wav_bytes) + kwargs: dict = {} + if language: + kwargs["language"] = language + result = model.transcribe(path, **kwargs) + text = (result.get("text") or "").strip() + return (text if text else None), None + except ImportError: + return None, "openai-whisper is not installed (pip install openai-whisper)" + except Exception as e: + return None, str(e)[:500] + finally: + if path: + try: + os.unlink(path) + except OSError: + pass + + +def compute_daily_omni_text_audio_consistency_metrics( + input_requests: list[Any], + outputs: list[Any], + *, + include_per_item: bool = False, +) -> dict[str, Any] | None: + """Compare option letter from ``generated_text`` vs Whisper transcript of output audio. + + Only considers requests where ``outputs[i]`` has ``generated_audio_wav_bytes`` set + (populated by the omni benchmark when TA check is enabled). + """ + if not input_requests or len(input_requests) != len(outputs): + return None + if not all(isinstance(r, DailyOmniSampleRequest) for r in input_requests): + return None + + ta_no_wav = 0 + ta_asr_failed = 0 + ta_text_unparsed = 0 + ta_audio_unparsed = 0 + ta_consistent = 0 + ta_mismatch = 0 + ta_both_parsed = 0 + items: list[dict[str, Any]] = [] + + for req, out in zip(input_requests, outputs, strict=True): + assert isinstance(req, DailyOmniSampleRequest) + rid = req.request_id + if not getattr(out, "success", False): + if include_per_item: + items.append( + { + "request_id": rid, + "skipped": True, + "reason": "request_not_success", + } + ) + continue + + wav = getattr(out, "generated_audio_wav_bytes", None) + if not wav: + ta_no_wav += 1 + if include_per_item: + items.append( + { + "request_id": rid, + "skipped": False, + "reason": "no_output_audio", + "text_choice": extract_predicted_choice(getattr(out, "generated_text", "") or ""), + } + ) + continue + + transcript, asr_err = transcribe_wav_bytes(wav) + if asr_err: + ta_asr_failed += 1 + if include_per_item: + items.append( + { + "request_id": rid, + "asr_error": asr_err, + "text_choice": extract_predicted_choice(getattr(out, "generated_text", "") or ""), + } + ) + continue + + text_choice = extract_predicted_choice(getattr(out, "generated_text", "") or "") + audio_choice = extract_choice_from_asr_transcript(transcript or "") + + if text_choice is None: + ta_text_unparsed += 1 + if audio_choice is None: + ta_audio_unparsed += 1 + + if text_choice is not None and audio_choice is not None: + ta_both_parsed += 1 + if text_choice == audio_choice: + ta_consistent += 1 + else: + ta_mismatch += 1 + + if include_per_item: + consistent: bool | None + if text_choice is None or audio_choice is None: + consistent = None + else: + consistent = text_choice == audio_choice + items.append( + { + "request_id": rid, + "text_choice": text_choice, + "audio_choice": audio_choice, + "asr_transcript": (transcript or "")[:500], + "text_audio_consistent": consistent, + } + ) + + comparable = ta_consistent + ta_mismatch + rate = (ta_consistent / comparable) if comparable else None + + out: dict[str, Any] = { + "daily_omni_ta_enabled": True, + "daily_omni_ta_no_output_audio": ta_no_wav, + "daily_omni_ta_asr_failed": ta_asr_failed, + "daily_omni_ta_text_unparsed": ta_text_unparsed, + "daily_omni_ta_audio_unparsed": ta_audio_unparsed, + "daily_omni_ta_both_parsed": ta_both_parsed, + "daily_omni_ta_consistent": ta_consistent, + "daily_omni_ta_mismatch": ta_mismatch, + "daily_omni_ta_consistency_rate": rate, + } + if include_per_item: + out["daily_omni_ta_items"] = items + return out + + +def print_daily_omni_text_audio_summary(metrics: dict[str, Any]) -> None: + if not metrics.get("daily_omni_ta_enabled"): + return + print("{s:{c}^{n}}".format(s=" Daily-Omni text vs audio (ASR) ", n=50, c="=")) + print("{:<40} {:<10}".format("No output audio captured:", metrics.get("daily_omni_ta_no_output_audio", 0))) + print("{:<40} {:<10}".format("ASR failed:", metrics.get("daily_omni_ta_asr_failed", 0))) + print("{:<40} {:<10}".format("Both text+audio letter parsed:", metrics.get("daily_omni_ta_both_parsed", 0))) + print("{:<40} {:<10}".format("Consistent (same letter):", metrics.get("daily_omni_ta_consistent", 0))) + print("{:<40} {:<10}".format("Mismatch:", metrics.get("daily_omni_ta_mismatch", 0))) + r = metrics.get("daily_omni_ta_consistency_rate") + if r is not None: + print("{:<40} {:<10.4f}".format("Consistency rate (of both parsed):", r)) + print( + "{:<40} {:<10}".format( + "Text unparsed (among w/ audio):", + metrics.get("daily_omni_ta_text_unparsed", 0), + ) + ) + print( + "{:<40} {:<10}".format( + "Audio unparsed (among w/ audio):", + metrics.get("daily_omni_ta_audio_unparsed", 0), + ) + ) diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py b/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py new file mode 100644 index 000000000..ca6de4cb2 --- /dev/null +++ b/vllm_omni/benchmarks/data_modules/seed_tts_dataset.py @@ -0,0 +1,272 @@ +"""Seed-TTS zero-shot evaluation-style prompts for ``vllm bench serve``. + +Loads rows from the `meta.lst` format used in `BytedanceSpeech/seed-tts-eval`_ (or any +HuggingFace dataset repo with the same layout):: + + utt_id|prompt_transcript|prompt_wav_relative_path|text_to_synthesize + +Each benchmark request supplies target text plus ``ref_text`` / ``ref_audio`` (Qwen3-TTS ``Base`` / +voice clone), merged into the JSON body. By default ``ref_audio`` is an inline ``data:`` URL so +the server does not need ``--allowed-local-media-path``. Use ``--seed-tts-file-ref-audio`` for +``file://`` (smaller bodies; requires that flag). Use ``--backend openai-audio-speech`` +(``/v1/audio/speech``) or ``--backend openai-chat-omni`` (``/v1/chat/completions`` with the same +fields on the body plus a Qwen3-Omni-style ``system`` message and the target text as ``user`` content). + +.. _BytedanceSpeech/seed-tts-eval: https://github.com/BytedanceSpeech/seed-tts-eval +""" + +from __future__ import annotations + +import base64 +import logging +import random +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from vllm.benchmarks.datasets import BenchmarkDataset, SampleRequest +from vllm.tokenizers import TokenizerLike +from vllm.tokenizers.hf import get_cached_tokenizer + +logger = logging.getLogger(__name__) + +# Matches Qwen3-Omni serving examples (``openai_chat_completion_client_for_multimodal_generation`` / +# ``qwen3_omni/gradio_demo``) plus explicit TTS / voice-clone instructions for chat completions. +SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT = ( + "You are Qwen, a virtual human developed by the Qwen Team, Alibaba Group, " + "capable of perceiving auditory and visual inputs, as well as generating text and speech.\n" + "For this request you act as a text-to-speech engine with zero-shot voice cloning: " + "the API provides reference audio and its transcript (ref_audio, ref_text) and task_type Base. " + "The user message is the exact text you must speak. " + "Synthesize natural speech in the same language as that user text, " + "matching the timbre, prosody, and speaking style of the reference audio while reading the new content clearly." +) + + +@dataclass +class SeedTTSSampleRequest(SampleRequest): + """``SampleRequest`` with per-row fields merged into ``/v1/audio/speech`` JSON.""" + + #: Shallow-merged into ``RequestFuncInput.extra_body`` (ref_audio, ref_text, task_type, …). + seed_tts_speech_extra: dict[str, Any] | None = None + seed_tts_utterance_id: str = "" + seed_tts_locale: str = "" + #: For ``openai-chat-omni``: becomes the chat ``system`` message (Qwen3-Omni + TTS behavior). + seed_tts_system_prompt: str = "" + #: Local path to reference prompt WAV (for SIM vs. synthesized PCM in ``seed_tts_eval``). + seed_tts_ref_wav_path: str = "" + + +@dataclass +class _SeedTTSRow: + utterance_id: str + ref_text: str + prompt_wav_rel: str + target_text: str + + +def _parse_meta_line(line: str) -> _SeedTTSRow | None: + line = line.strip() + if not line or line.startswith("#"): + return None + parts = line.split("|") + if len(parts) < 4: + logger.warning("Skipping malformed meta.lst line (need 4 '|'-fields): %r", line[:120]) + return None + utt_id, ref_text, wav_rel, target = parts[0], parts[1], parts[2], parts[3] + if not target.strip(): + return None + return _SeedTTSRow( + utterance_id=utt_id.strip(), + ref_text=ref_text.strip(), + prompt_wav_rel=wav_rel.strip(), + target_text=target.strip(), + ) + + +def _load_meta_rows(meta_file: Path) -> list[_SeedTTSRow]: + text = meta_file.read_text(encoding="utf-8") + rows: list[_SeedTTSRow] = [] + for line in text.splitlines(): + r = _parse_meta_line(line) + if r is not None: + rows.append(r) + return rows + + +def resolve_seed_tts_root(dataset_path: str | None, *, explicit_root: str | None) -> Path: + """Return directory containing ``{locale}/meta.lst`` and ``{locale}/prompt-wavs/``.""" + if explicit_root: + root = Path(explicit_root).expanduser().resolve() + if not root.is_dir(): + raise FileNotFoundError(f"--seed-tts-root is not a directory: {root}") + return root + + if not dataset_path: + raise ValueError("Seed-TTS requires --dataset-path (HF repo id or local root) or --seed-tts-root.") + + p = Path(dataset_path).expanduser() + if p.exists() and p.is_dir(): + return p.resolve() + + repo_id = dataset_path.strip() + try: + from huggingface_hub import snapshot_download + except ImportError as e: + raise ImportError( + "Install huggingface_hub to download Seed-TTS from the Hub, or clone the dataset " + "locally and pass --dataset-path / --seed-tts-root to that directory." + ) from e + cache = snapshot_download(repo_id=repo_id, repo_type="dataset") + return Path(cache).resolve() + + +def _ref_audio_payload(wav_path: Path, *, inline: bool) -> str: + if inline: + raw = wav_path.read_bytes() + b64 = base64.b64encode(raw).decode("ascii") + return f"data:audio/wav;base64,{b64}" + return wav_path.expanduser().resolve().as_uri() + + +class SeedTTSDataset(BenchmarkDataset): + """Seed-TTS-style zero-shot TTS rows for throughput/latency benchmarking. + + Args: + dataset_path: HuggingFace dataset repo id (``org/dataset``) or local directory with + ``en/meta.lst`` (and ``zh/meta.lst`` if using zh). + locale: ``en`` or ``zh`` — which subfolder under the root to read. + inline_ref_audio: If True (default), embed prompt WAV as ``data:audio/wav;base64,...`` + so Qwen3-TTS / ``/v1/audio/speech`` works without server + ``--allowed-local-media-path``. If False, use ``file://`` (smaller + requests; server must set ``--allowed-local-media-path`` to the dataset root). + seed_tts_root: Optional override for the root directory (same layout as HF dataset). + system_prompt: Optional override for the chat system message when using + ``--backend openai-chat-omni``; defaults to :data:`SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT`. + """ + + IS_MULTIMODAL = False + DEFAULT_OUTPUT_LEN = 2048 + + def __init__( + self, + dataset_path: str, + random_seed: int = 0, + locale: str = "en", + inline_ref_audio: bool = True, + seed_tts_root: str | None = None, + system_prompt: str | None = None, + disable_shuffle: bool = False, + **kwargs: Any, + ) -> None: + if locale not in ("en", "zh"): + raise ValueError("locale must be 'en' or 'zh'") + self.locale = locale + self.inline_ref_audio = inline_ref_audio + self._explicit_root = seed_tts_root + sp = (system_prompt or "").strip() + self._system_prompt = sp if sp else SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT + super().__init__( + dataset_path=dataset_path, + random_seed=random_seed, + disable_shuffle=disable_shuffle, + **kwargs, + ) + self._root = resolve_seed_tts_root(self.dataset_path, explicit_root=self._explicit_root) + self._rows: list[_SeedTTSRow] = [] + self.load_data() + + def load_data(self) -> None: + meta = self._root / self.locale / "meta.lst" + if not meta.is_file(): + raise FileNotFoundError( + f"Seed-TTS meta not found: {meta}. " + f"Expected layout from seed-tts-eval (e.g. {self._root}/{self.locale}/meta.lst)." + ) + self._rows = _load_meta_rows(meta) + if not self._rows: + raise ValueError(f"No valid rows in {meta}") + if not self.disable_shuffle: + rng = random.Random(self.random_seed) + rng.shuffle(self._rows) + self.data = self._rows + logger.info( + "Loaded Seed-TTS: root=%s locale=%s rows=%d inline_ref_audio=%s", + self._root, + self.locale, + len(self._rows), + self.inline_ref_audio, + ) + + def sample( + self, + tokenizer: TokenizerLike, + num_requests: int, + output_len: int | None = None, + request_id_prefix: str = "", + no_oversample: bool = False, + **kwargs: Any, + ) -> list[SampleRequest]: + if output_len is None: + output_len = self.DEFAULT_OUTPUT_LEN + + tok = get_cached_tokenizer(tokenizer) + out: list[SampleRequest] = [] + for i, row in enumerate(self._rows): + if len(out) >= num_requests: + break + wav_path = (self._root / self.locale / row.prompt_wav_rel).resolve() + if not wav_path.is_file(): + logger.warning("Missing prompt wav for %s: %s", row.utterance_id, wav_path) + continue + + target = row.target_text + prompt_len = len(tok.encode(f"{self._system_prompt}\n{target}")) + lang = "English" if self.locale == "en" else "Chinese" + ref_uri = _ref_audio_payload(wav_path, inline=self.inline_ref_audio) + speech_extra: dict[str, Any] = { + "ref_audio": ref_uri, + "ref_text": row.ref_text, + "task_type": "Base", + "language": lang, + "max_new_tokens": output_len, + } + + out.append( + SeedTTSSampleRequest( + prompt=target, + prompt_len=prompt_len, + expected_output_len=output_len, + multi_modal_data=None, + request_id=f"{request_id_prefix}{i}", + seed_tts_speech_extra=speech_extra, + seed_tts_utterance_id=row.utterance_id, + seed_tts_locale=self.locale, + seed_tts_system_prompt=self._system_prompt, + seed_tts_ref_wav_path=str(wav_path), + ) + ) + + logger.info("Seed-TTS: built %d requests (asked %d)", len(out), num_requests) + self.maybe_oversample_requests(out, num_requests, request_id_prefix, no_oversample) + return out + + +def load_seed_tts_dataset( + dataset_path: str, + random_seed: int = 0, + locale: str = "en", + inline_ref_audio: bool = True, + seed_tts_root: str | None = None, + system_prompt: str | None = None, + **kwargs: Any, +) -> SeedTTSDataset: + return SeedTTSDataset( + dataset_path=dataset_path, + random_seed=random_seed, + locale=locale, + inline_ref_audio=inline_ref_audio, + seed_tts_root=seed_tts_root, + system_prompt=system_prompt, + **kwargs, + ) diff --git a/vllm_omni/benchmarks/data_modules/seed_tts_eval.py b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py new file mode 100644 index 000000000..d5f1b6470 --- /dev/null +++ b/vllm_omni/benchmarks/data_modules/seed_tts_eval.py @@ -0,0 +1,729 @@ +"""Seed-TTS WER aligned with Bytedance ``seed-tts-eval`` / ``run_wer.py``. + +Matches the published protocol (see Hugging Face dataset card and +https://github.com/BytedanceSpeech/seed-tts-eval): + +- **EN**: ``openai/whisper-large-v3`` via ``transformers``, audio resampled to **16 kHz** + (same as ``run_wer.py``). +- **ZH**: ``funasr`` **paraformer-zh**, hypothesis converted with **zhconv** to zh-cn. +- **WER**: ``jiwer`` after punctuation stripping (``zhon.hanzi.punctuation`` + ``string.punctuation``, + preserving ``'``) and EN lowercasing / ZH per-character spacing. Supports jiwer 3.x + (``compute_measures``) and 4.x (``process_words``). + +- **SIM** (speaker similarity proxy): cosine similarity of L2-normalized mean-pooled **WavLM** + embeddings (reference prompt WAV vs. synthesized PCM), 16 kHz. Official ``cal_sim.sh`` uses + UniSpeech ``verification_pair_list_v2.py`` with a **fine-tuned** WavLM SV checkpoint — set + ``SEED_TTS_WAVLM_MODEL`` to another HF id if you need closer parity. Disable with + ``SEED_TTS_SIM_EVAL=0``. Optional: ``SEED_TTS_SIM_DEVICE`` (e.g. ``cpu``) to avoid GPU + issues when Whisper already uses CUDA; ``SEED_TTS_WAVLM_MIN_SAMPLES`` pads very short + waveforms so the WavLM CNN front-end does not fail. + +- **UTMOS** (predicted MOS from TorchScript): default ``balacoon/utmos`` → ``utmos.jit`` + (Sarulab-style demo export). Uses ``torch`` + ``huggingface_hub`` only. Aggregate metrics + are over **all requests with captured PCM** (independent of ASR/WER). Non-finite scores are + dropped and counted as failures. Override repo/file via ``SEED_TTS_UTMOS_HF_REPO`` / + ``SEED_TTS_UTMOS_JIT_FILE``. **Device**: defaults to **CPU** when ``SEED_TTS_UTMOS_DEVICE`` + is unset; set ``SEED_TTS_UTMOS_DEVICE=cuda:0`` (or ``cuda:1`` etc.) to run on GPU. The JIT + model is loaded directly onto the target device via ``map_location`` to avoid cross-device + issues (some PyTorch builds/Windows have problems moving TorchScript modules after load). + Forward uses **float32** waveform in ``[-1, 1]`` (same as the WER resampled array) so + tensor dtypes match JIT weights; using int16 triggers + ``RuntimeError: input type and weight type should be same`` on common exports. Disable + with ``SEED_TTS_UTMOS_EVAL=0``. + +Enable with ``SEED_TTS_WER_EVAL=1`` or ``--seed-tts-wer-eval``. Install optional deps:: + + pip install 'vllm-omni[seed-tts-eval]' + +Env: ``SEED_TTS_EVAL_DEVICE`` (e.g. ``cuda:0``, ``cpu``); ``SEED_TTS_HF_WHISPER_MODEL`` +defaults to ``openai/whisper-large-v3`` (override for debugging only). +""" + +from __future__ import annotations + +import io +import logging +import math +import os +import statistics +import string +import tempfile +import threading +import wave +from typing import Any + +import numpy as np +from vllm.benchmarks.datasets import SampleRequest + +from vllm_omni.benchmarks.data_modules.seed_tts_dataset import SeedTTSSampleRequest + +logger = logging.getLogger(__name__) + +# Mirrors seed-tts-eval/run_wer.py +OFFICIAL_WHISPER_HF_ID = "openai/whisper-large-v3" +PARAFORMER_MODEL_ID = "paraformer-zh" + +_lock = threading.Lock() +_device: str | None = None +_en_processor = None +_en_model = None +_zh_paraformer = None +_wavlm_model = None +_wavlm_processor = None +_wavlm_device: str | None = None +_utmos_jit_model = None +_utmos_jit_device: str | None = None +_utmos_jit_load_failed = False +_utmos_forward_warned = False + + +def pcm_s16le_mono_to_wav_bytes(pcm: bytes, *, sample_rate: int = 24000) -> bytes: + buf = io.BytesIO() + with wave.open(buf, "wb") as wf: + wf.setnchannels(1) + wf.setsampwidth(2) + wf.setframerate(sample_rate) + wf.writeframes(pcm) + return buf.getvalue() + + +def _get_eval_device() -> str: + explicit = os.environ.get("SEED_TTS_EVAL_DEVICE", "").strip() + if explicit: + return explicit + try: + import torch + + return "cuda:0" if torch.cuda.is_available() else "cpu" + except ImportError: + return "cpu" + + +def _punctuation_all() -> str: + from zhon.hanzi import punctuation + + return punctuation + string.punctuation + + +def _jiwer_wer(reference: str, hypothesis: str) -> float: + """Word-level WER; strings are normalized like ``run_wer.process_one``. + + jiwer 4.x removed ``compute_measures`` (``ImportError``); fall back to ``process_words``. + """ + try: + from jiwer import compute_measures + + return float(compute_measures(reference, hypothesis)["wer"]) + except ImportError: + import jiwer + + out = jiwer.process_words(reference, hypothesis) + return float(out.wer) + + +def process_one_official(hypo: str, truth: str, lang: str) -> tuple[float, str, str]: + """Same normalization + ``jiwer`` call as ``run_wer.process_one`` (hypo=ASR, truth=reference).""" + raw_truth = truth + raw_hypo = hypo + truth_n = truth + hypo_n = hypo + for x in _punctuation_all(): + if x == "'": + continue + truth_n = truth_n.replace(x, "") + hypo_n = hypo_n.replace(x, "") + truth_n = truth_n.replace(" ", " ") + hypo_n = hypo_n.replace(" ", " ") + if lang == "zh": + truth_n = " ".join([x for x in truth_n]) + hypo_n = " ".join([x for x in hypo_n]) + elif lang == "en": + truth_n = truth_n.lower() + hypo_n = hypo_n.lower() + else: + raise ValueError(f"unsupported lang {lang!r}") + wer = _jiwer_wer(truth_n, hypo_n) + return wer, raw_truth, raw_hypo + + +def _pcm_s16le_to_f32_16k(pcm: bytes, pcm_sample_rate: int = 24000) -> np.ndarray: + import scipy.signal + + if not pcm: + return np.zeros(0, dtype=np.float32) + raw = np.frombuffer(pcm, dtype=np.int16).astype(np.float32) / 32768.0 + target_len = int(len(raw) * 16000 / pcm_sample_rate) + if target_len <= 0: + return np.zeros(0, dtype=np.float32) + return scipy.signal.resample(raw, target_len).astype(np.float32) + + +def _eval_submetric_enabled(env_name: str, *, default: bool = True) -> bool: + raw = os.environ.get(env_name, "").strip().lower() + if raw in ("0", "false", "no", "off"): + return False + if raw in ("1", "true", "yes", "on"): + return True + return default + + +def _audio_path_to_f32_16k(path: str) -> np.ndarray: + import scipy.signal + import soundfile as sf + + data, sr = sf.read(path, dtype="float32", always_2d=True) + mono = np.mean(data, axis=1).astype(np.float32) + if int(sr) == 16000: + return mono + target_len = max(1, int(len(mono) * 16000 / int(sr))) + return scipy.signal.resample(mono, target_len).astype(np.float32) + + +def _ensure_wavlm_sim() -> None: + global _wavlm_model, _wavlm_processor, _wavlm_device + with _lock: + if _wavlm_model is not None: + return + from transformers import AutoFeatureExtractor, AutoModel + + mid = os.environ.get("SEED_TTS_WAVLM_MODEL", "microsoft/wavlm-base-plus").strip() or "microsoft/wavlm-base-plus" + _wavlm_device = os.environ.get("SEED_TTS_SIM_DEVICE", "").strip() or _get_eval_device() + logger.warning( + "Loading WavLM %r on %s for Seed-TTS SIM (embedding cosine; not identical to " + "seed-tts-eval UniSpeech SV checkpoint).", + mid, + _wavlm_device, + ) + _wavlm_processor = AutoFeatureExtractor.from_pretrained(mid) + _wavlm_model = AutoModel.from_pretrained(mid).to(_wavlm_device) + _wavlm_model.eval() + + +def _wavlm_prepare_waveform(wav: np.ndarray) -> np.ndarray: + """Trim, pad to a minimum length WavLM/Wav2Vec2 CNN stack accepts, float32 mono.""" + max_sec = float(os.environ.get("SEED_TTS_WAVLM_MAX_SECONDS", "30")) + cap = int(max_sec * 16000) + w = np.asarray(wav, dtype=np.float32).reshape(-1) + if len(w) == 0: + return w + if len(w) > cap: + w = w[:cap].copy() + # Very short clips make the strided conv front-end fail (shape / empty time dim). + min_samples = int(os.environ.get("SEED_TTS_WAVLM_MIN_SAMPLES", "4000")) + if len(w) < min_samples: + w = np.pad(w, (0, min_samples - len(w)), mode="constant") + return w + + +def _wavlm_mean_embedding_f32_16k(wav: np.ndarray) -> np.ndarray | None: + import torch + + _ensure_wavlm_sim() + w = _wavlm_prepare_waveform(wav) + if len(w) == 0: + return None + assert _wavlm_processor is not None and _wavlm_model is not None and _wavlm_device is not None + # Single utterance: avoid padding=True (adds zeros that distort mean pooling). Still pass + # attention_mask when the extractor provides it (sample-level; do not mix with hidden length). + try: + inputs = _wavlm_processor( + w, + sampling_rate=16000, + return_tensors="pt", + padding=False, + return_attention_mask=True, + ) + except TypeError: + inputs = _wavlm_processor( + w, + sampling_rate=16000, + return_tensors="pt", + padding=False, + ) + iv = inputs["input_values"].to(_wavlm_device) + am = inputs.get("attention_mask") + if am is not None: + am = am.to(_wavlm_device) + with torch.inference_mode(): + out = _wavlm_model(iv, attention_mask=am) + h = out.last_hidden_state + v = h.mean(dim=1).squeeze(0).float().cpu().numpy() + n = float(np.linalg.norm(v)) + if not np.isfinite(n) or n < 1e-8: + return None + return (v / n).astype(np.float32) + + +def _cosine_similarity_unit_vectors(a: np.ndarray, b: np.ndarray) -> float: + return float(np.dot(a, b)) + + +def _ensure_utmos_jit_model() -> Any | None: + """Load UTMOS as TorchScript (``balacoon/utmos`` style): no ``import utmos`` / fairseq.""" + global _utmos_jit_model, _utmos_jit_device, _utmos_jit_load_failed + with _lock: + if _utmos_jit_load_failed: + return None + if _utmos_jit_model is not None: + return _utmos_jit_model + try: + import torch + from huggingface_hub import hf_hub_download + + repo = os.environ.get("SEED_TTS_UTMOS_HF_REPO", "balacoon/utmos").strip() or "balacoon/utmos" + fname = os.environ.get("SEED_TTS_UTMOS_JIT_FILE", "utmos.jit").strip() or "utmos.jit" + logger.warning( + "Loading UTMOS TorchScript from Hugging Face %r file %r (one-time download/cache)...", + repo, + fname, + ) + path = hf_hub_download(repo_id=repo, filename=fname, repo_type="model") + + # TODO The model weights in UTMOS must be loaded in cuda:0; otherwise, the model execution will fail. + want = "cuda:0" + if want.startswith("cuda") and torch.cuda.is_available(): + idx = want.split(":")[-1] if ":" in want else "0" + target_dev = f"cuda:{idx}" + else: + target_dev = "cpu" + + try: + m = torch.jit.load(path, map_location=target_dev) + m.eval() + _utmos_jit_device = target_dev + except Exception as load_e: + if target_dev.startswith("cuda"): + logger.warning( + "UTMOS JIT load on %s failed (%s), retrying on CPU...", + target_dev, + load_e, + ) + m = torch.jit.load(path, map_location="cpu") + m.eval() + _utmos_jit_device = "cpu" + else: + raise + _utmos_jit_model = m + except Exception as e: + logger.warning( + "UTMOS JIT unavailable (install torch + huggingface_hub; check HF access): %s", + e, + ) + _utmos_jit_load_failed = True + return None + return _utmos_jit_model + + +def _utmos_predict_f32_16k(wav_f32: np.ndarray) -> float | None: + """MOS from JIT model; input is float32 mono @ 16 kHz in ``[-1, 1]`` (WER pipeline). + + ``balacoon/utmos`` demos sometimes use int16 numpy, but the exported ``.jit`` weights are + float32; passing int16 tensors causes: "RuntimeError: ... input type and weight type + should be same". + """ + import torch + + if len(wav_f32) == 0: + return None + model = _ensure_utmos_jit_model() + if model is None: + return None + # Infer model's device from its first parameter/buffer to guarantee input sits with weights. + try: + model_dev = next(model.parameters()).device + except StopIteration: + try: + model_dev = next(model.buffers()).device + except StopIteration: + model_dev = torch.device("cpu") + w = np.ascontiguousarray(wav_f32, dtype=np.float32) + x = torch.from_numpy(w).unsqueeze(0).to(device=model_dev, dtype=torch.float32) + with torch.no_grad(): + out = model(x) + val = float(out.reshape(-1)[0].item()) + if not math.isfinite(val): + return None + return val + + +def _ensure_en_asr() -> None: + global _en_processor, _en_model, _device + with _lock: + if _en_processor is not None: + return + from transformers import WhisperForConditionalGeneration, WhisperProcessor + + _device = _get_eval_device() + mid = os.environ.get("SEED_TTS_HF_WHISPER_MODEL", OFFICIAL_WHISPER_HF_ID).strip() or OFFICIAL_WHISPER_HF_ID + logger.warning( + "Loading Seed-TTS eval Whisper HF model %r on %s (one-time, seed-tts-eval protocol)...", + mid, + _device, + ) + _en_processor = WhisperProcessor.from_pretrained(mid) + _en_model = WhisperForConditionalGeneration.from_pretrained(mid).to(_device) + _en_model.eval() + + +def _ensure_zh_asr() -> None: + global _zh_paraformer, _device + with _lock: + if _zh_paraformer is not None: + return + from funasr import AutoModel + + _device = _get_eval_device() + logger.warning( + "Loading Seed-TTS eval Paraformer %r on %s (one-time, seed-tts-eval protocol)...", + PARAFORMER_MODEL_ID, + _device, + ) + try: + _zh_paraformer = AutoModel(model=PARAFORMER_MODEL_ID, device=_device) + except TypeError: + _zh_paraformer = AutoModel(model=PARAFORMER_MODEL_ID) + + +def _transcribe_en_f32_16k(wav_f32: np.ndarray) -> str: + import torch + + _ensure_en_asr() + if len(wav_f32) == 0: + return "" + with _lock: + assert _en_processor is not None and _en_model is not None and _device is not None + inputs = _en_processor(wav_f32, sampling_rate=16000, return_tensors="pt") + input_features = inputs.input_features.to(_device) + with torch.no_grad(): + try: + forced = _en_processor.get_decoder_prompt_ids(language="english", task="transcribe") + predicted_ids = _en_model.generate(input_features, forced_decoder_ids=forced) + except Exception: + predicted_ids = _en_model.generate( + input_features, + language="english", + task="transcribe", + ) + text = _en_processor.batch_decode(predicted_ids, skip_special_tokens=True)[0] + return (text or "").strip() + + +def _transcribe_zh_wav_path(wav_path: str) -> str: + import zhconv + + _ensure_zh_asr() + with _lock: + assert _zh_paraformer is not None + res = _zh_paraformer.generate(input=wav_path, batch_size_s=300) + transcription = res[0]["text"] if res else "" + return zhconv.convert(transcription, "zh-cn").strip() + + +def _missing_deps_message(lang: str) -> str | None: + try: + import jiwer # noqa: F401 + from zhon.hanzi import punctuation # noqa: F401 + except ImportError as e: + return f"Seed-TTS WER eval needs jiwer and zhon ({e!s}). Install: pip install 'vllm-omni[seed-tts-eval]'" + try: + import scipy.signal # noqa: F401 + import soundfile # noqa: F401 + except ImportError as e: + return f"Seed-TTS WER eval needs scipy and soundfile ({e!s})." + if lang == "en": + try: + import torch # noqa: F401 + from transformers import WhisperForConditionalGeneration # noqa: F401 + except ImportError as e: + return f"English WER needs torch and transformers ({e!s}). Install: pip install 'vllm-omni[seed-tts-eval]'" + else: + try: + import zhconv # noqa: F401 + from funasr import AutoModel # noqa: F401 + except ImportError as e: + return f"Chinese WER needs funasr and zhconv ({e!s}). Install: pip install 'vllm-omni[seed-tts-eval]'" + return None + + +def compute_seed_tts_wer_metrics( + input_requests: list[SampleRequest], + outputs: list[Any], + *, + include_per_item: bool = False, +) -> dict[str, Any] | None: + """If all requests are :class:`SeedTTSSampleRequest`, run seed-tts-eval-style WER.""" + global _utmos_forward_warned + if not input_requests or len(input_requests) != len(outputs): + return None + if not all(isinstance(r, SeedTTSSampleRequest) for r in input_requests): + return None + + first = input_requests[0] + assert isinstance(first, SeedTTSSampleRequest) + lang = "zh" if (first.seed_tts_locale or "en").lower().startswith("zh") else "en" + + setup_err = _missing_deps_message(lang) + if setup_err: + logger.error("%s", setup_err) + return { + "seed_tts_eval_setup_error": setup_err, + "seed_tts_eval_protocol": "seed-tts-eval", + "seed_tts_content_evaluated": 0, + "seed_tts_content_error_mean": None, + "seed_tts_content_error_median": None, + "seed_tts_request_failed": 0, + "seed_tts_no_pcm": 0, + "seed_tts_asr_failed": 0, + "seed_tts_content_metric": "wer", + } + + import soundfile as sf + + errs: list[float] = [] + items: list[dict[str, Any]] = [] + asr_failed = 0 + no_pcm = 0 + request_failed = 0 + sim_values: list[float] = [] + utmos_values: list[float] = [] + sim_failed = 0 + sim_skipped_no_ref = 0 + utmos_failed = 0 + utmos_on = _eval_submetric_enabled("SEED_TTS_UTMOS_EVAL", default=True) + + for req, out in zip(input_requests, outputs, strict=True): + assert isinstance(req, SeedTTSSampleRequest) + ref = req.prompt + locale = req.seed_tts_locale or "en" + row_lang = "zh" if locale.lower().startswith("zh") else "en" + utmos_v: float | None = None + + if not out.success: + request_failed += 1 + if include_per_item: + items.append( + { + "utterance_id": req.seed_tts_utterance_id, + "locale": locale, + "error": "request_failed", + "detail": (out.error or "")[:500], + } + ) + continue + + pcm = getattr(out, "tts_output_pcm_bytes", None) + if not pcm: + no_pcm += 1 + if include_per_item: + items.append( + { + "utterance_id": req.seed_tts_utterance_id, + "locale": locale, + "error": "no_pcm", + } + ) + continue + + wav_16k = _pcm_s16le_to_f32_16k(pcm) + if len(wav_16k) == 0: + asr_failed += 1 + if include_per_item: + items.append( + { + "utterance_id": req.seed_tts_utterance_id, + "locale": locale, + "error": "empty_audio", + } + ) + continue + + # UTMOS scores synthesized audio only; do not gate on ASR/WER (those can fail independently). + if utmos_on: + try: + utmos_v = _utmos_predict_f32_16k(wav_16k) + if utmos_v is not None: + utmos_values.append(utmos_v) + elif not _utmos_jit_load_failed: + utmos_failed += 1 + except Exception: + if not _utmos_forward_warned: + _utmos_forward_warned = True + logger.warning( + "UTMOS JIT forward failed (first utterance=%s; set logging DEBUG for " + "full trace). Check sample rate (16 kHz), input shape, or " + "SEED_TTS_UTMOS_DEVICE.", + req.seed_tts_utterance_id, + exc_info=True, + ) + else: + logger.debug( + "UTMOS forward failed for %s", + req.seed_tts_utterance_id, + exc_info=True, + ) + utmos_failed += 1 + + try: + if row_lang == "en": + hyp = _transcribe_en_f32_16k(wav_16k) + else: + fd, tmp_wav = tempfile.mkstemp(suffix=".wav") + os.close(fd) + try: + sf.write(tmp_wav, wav_16k, 16000, subtype="PCM_16") + hyp = _transcribe_zh_wav_path(tmp_wav) + finally: + try: + os.unlink(tmp_wav) + except OSError: + pass + except Exception as e: + logger.exception("Seed-TTS ASR failed for %s", req.seed_tts_utterance_id) + asr_failed += 1 + if include_per_item: + items.append( + { + "utterance_id": req.seed_tts_utterance_id, + "locale": locale, + "error": "asr_exception", + "detail": str(e)[:500], + } + ) + continue + + if not hyp: + asr_failed += 1 + if include_per_item: + items.append( + { + "utterance_id": req.seed_tts_utterance_id, + "locale": locale, + "error": "empty_asr", + } + ) + continue + + try: + wer, raw_truth, raw_hypo = process_one_official(hyp, ref, row_lang) + except Exception as e: + logger.warning("jiwer/normalize failed for %s: %s", req.seed_tts_utterance_id, e) + asr_failed += 1 + if include_per_item: + items.append( + { + "utterance_id": req.seed_tts_utterance_id, + "locale": locale, + "error": "wer_compute_failed", + "detail": str(e)[:500], + } + ) + continue + + errs.append(wer) + sim_v: float | None = None + + if _eval_submetric_enabled("SEED_TTS_SIM_EVAL", default=True): + ref_path = getattr(req, "seed_tts_ref_wav_path", "") or "" + if ref_path and os.path.isfile(ref_path): + try: + ref_wav = _audio_path_to_f32_16k(ref_path) + e_ref = _wavlm_mean_embedding_f32_16k(ref_wav) + e_hyp = _wavlm_mean_embedding_f32_16k(wav_16k) + if e_ref is not None and e_hyp is not None: + sim_v = _cosine_similarity_unit_vectors(e_ref, e_hyp) + sim_values.append(sim_v) + except Exception as e: + logger.warning( + "SIM embedding failed for utterance=%s: %s: %s", + req.seed_tts_utterance_id, + type(e).__name__, + e, + ) + sim_failed += 1 + else: + sim_skipped_no_ref += 1 + + if include_per_item: + row: dict[str, Any] = { + "utterance_id": req.seed_tts_utterance_id, + "locale": locale, + "wer": wer, + "reference_raw": raw_truth, + "asr_raw": raw_hypo, + } + if sim_v is not None: + row["sim"] = sim_v + if utmos_v is not None: + row["utmos"] = utmos_v + items.append(row) + + result: dict[str, Any] = { + "seed_tts_eval_protocol": "seed-tts-eval", + "seed_tts_content_evaluated": len(errs), + "seed_tts_content_error_mean": statistics.fmean(errs) if errs else None, + "seed_tts_content_error_median": statistics.median(errs) if errs else None, + "seed_tts_request_failed": request_failed, + "seed_tts_no_pcm": no_pcm, + "seed_tts_asr_failed": asr_failed, + "seed_tts_content_metric": "wer", + "seed_tts_sim_evaluated": len(sim_values), + "seed_tts_sim_mean": statistics.fmean(sim_values) if sim_values else None, + "seed_tts_sim_median": statistics.median(sim_values) if sim_values else None, + "seed_tts_sim_failed": sim_failed, + "seed_tts_sim_skipped_no_ref": sim_skipped_no_ref, + "seed_tts_utmos_evaluated": len(utmos_values), + "seed_tts_utmos_mean": statistics.fmean(utmos_values) if utmos_values else None, + "seed_tts_utmos_median": statistics.median(utmos_values) if utmos_values else None, + "seed_tts_utmos_failed": utmos_failed, + } + if include_per_item: + result["seed_tts_wer_eval_items"] = items + return result + + +def print_seed_tts_wer_summary(metrics: dict[str, Any]) -> None: + setup = metrics.get("seed_tts_eval_setup_error") + if setup: + print("{s:{c}^{n}}".format(s=" Seed-TTS eval (seed-tts-eval protocol) ", n=50, c="=")) + print(setup) + return + + ev = int(metrics.get("seed_tts_content_evaluated", 0) or 0) + rf = int(metrics.get("seed_tts_request_failed", 0) or 0) + npc = int(metrics.get("seed_tts_no_pcm", 0) or 0) + af = int(metrics.get("seed_tts_asr_failed", 0) or 0) + sim_ev = int(metrics.get("seed_tts_sim_evaluated", 0) or 0) + ut_ev = int(metrics.get("seed_tts_utmos_evaluated", 0) or 0) + if ev == 0 and rf == 0 and npc == 0 and af == 0 and sim_ev == 0 and ut_ev == 0: + return + print("{s:{c}^{n}}".format(s=" Seed-TTS eval (seed-tts-eval protocol) ", n=50, c="=")) + print("{:<40} {:<10}".format("Evaluated (WER, lower is better):", ev)) + mean = metrics.get("seed_tts_content_error_mean") + if mean is not None: + print("{:<40} {:<10.4f}".format("Mean WER:", float(mean))) + med = metrics.get("seed_tts_content_error_median") + if med is not None: + print("{:<40} {:<10.4f}".format("Median WER:", float(med))) + print("{:<40} {:<10}".format("Request failed:", metrics.get("seed_tts_request_failed", 0))) + print("{:<40} {:<10}".format("No PCM captured:", metrics.get("seed_tts_no_pcm", 0))) + print("{:<40} {:<10}".format("ASR / WER failed:", metrics.get("seed_tts_asr_failed", 0))) + if sim_ev or metrics.get("seed_tts_sim_skipped_no_ref") or metrics.get("seed_tts_sim_failed"): + print("{:<40} {:<10}".format("SIM evaluated (higher ~ closer):", sim_ev)) + sm = metrics.get("seed_tts_sim_mean") + if sm is not None: + print("{:<40} {:<10.4f}".format("Mean SIM:", float(sm))) + s_med = metrics.get("seed_tts_sim_median") + if s_med is not None: + print("{:<40} {:<10.4f}".format("Median SIM:", float(s_med))) + print("{:<40} {:<10}".format("SIM skipped (no ref path):", metrics.get("seed_tts_sim_skipped_no_ref", 0))) + print("{:<40} {:<10}".format("SIM embedding errors:", metrics.get("seed_tts_sim_failed", 0))) + if ut_ev or metrics.get("seed_tts_utmos_failed"): + print("{:<40} {:<10}".format("UTMOS evaluated (JIT MOS, higher better):", ut_ev)) + um = metrics.get("seed_tts_utmos_mean") + if um is not None: + print("{:<40} {:<10.4f}".format("Mean UTMOS:", float(um))) + u_med = metrics.get("seed_tts_utmos_median") + if u_med is not None: + print("{:<40} {:<10.4f}".format("Median UTMOS:", float(u_med))) + print("{:<40} {:<10}".format("UTMOS errors:", metrics.get("seed_tts_utmos_failed", 0))) + print("=" * 50) diff --git a/vllm_omni/benchmarks/patch/__init__.py b/vllm_omni/benchmarks/patch/__init__.py index e69de29bb..ca6b41ba8 100644 --- a/vllm_omni/benchmarks/patch/__init__.py +++ b/vllm_omni/benchmarks/patch/__init__.py @@ -0,0 +1,3 @@ +"""Omni benchmark monkey-patches (side effects in ``patch.patch``).""" + +from . import patch as _patch_module # noqa: F401 diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 17d7498ba..41aed0942 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -6,6 +6,7 @@ import json import os import random import ssl +import sys import time import traceback from collections.abc import Iterable @@ -33,15 +34,245 @@ from vllm.logger import init_logger from vllm.tokenizers import TokenizerLike logger = init_logger(__name__) + +from vllm_omni.benchmarks.data_modules.daily_omni_dataset import DailyOmniDataset, DailyOmniSampleRequest from vllm_omni.benchmarks.data_modules.random_multi_modal_dataset import OmniRandomMultiModalDataset +from vllm_omni.benchmarks.data_modules.seed_tts_dataset import ( + SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT, + SeedTTSDataset, + SeedTTSSampleRequest, +) get_samples_old = datasets.get_samples +_DEFAULT_DAILY_OMNI_REPO = "liarliar/Daily-Omni" + + +def _seed_tts_capture_pcm_for_wer() -> bool: + return os.environ.get("SEED_TTS_WER_EVAL", "").lower() in ( + "1", + "true", + "yes", + ) + + +def _merge_extra_body_mm_kwargs(base: dict | None, overlay: dict | None) -> dict | None: + """Shallow-merge ``extra_body`` dicts; deep-merge ``mm_processor_kwargs`` if both set.""" + if not base and not overlay: + return None + out = dict(base or {}) + if not overlay: + return out + for k, v in overlay.items(): + if k == "mm_processor_kwargs" and isinstance(v, dict): + prev = out.get("mm_processor_kwargs") + merged_kw = {**(prev if isinstance(prev, dict) else {}), **v} + out["mm_processor_kwargs"] = merged_kw + else: + out[k] = v + return out + + +def _attach_daily_omni_to_request_func_input(sample: SampleRequest, rfi: RequestFuncInput) -> None: + """Apply per-request OpenAI fields (``mm_processor_kwargs``, messages) for Daily-Omni.""" + if not isinstance(sample, DailyOmniSampleRequest): + return + rfi.extra_body = _merge_extra_body_mm_kwargs(rfi.extra_body, sample.omni_extra_body) + if sample.omni_chat_messages is not None: + setattr(rfi, "omni_chat_messages", sample.omni_chat_messages) + else: + setattr(rfi, "mm_position", sample.omni_chat_mm_position) + + +def _attach_seed_tts_to_request_func_input(sample: SampleRequest, rfi: RequestFuncInput) -> None: + """Merge Seed-TTS per-row TTS fields (ref_audio, ref_text, task_type, …) into ``extra_body``. + + Used by both ``/v1/audio/speech`` and ``/v1/chat/completions`` (flattened into JSON body). + For ``openai-chat-omni``, also sets ``omni_chat_messages`` (system + user) so Qwen3-Omni + follows the same role layout as official TTS / multimodal demos. ``/v1/audio/speech`` ignores + ``messages`` and only uses ``input`` + body fields. + Flags ``openai-chat-omni`` to request audio output and optionally export PCM for WER. + """ + if not isinstance(sample, SeedTTSSampleRequest): + return + ex = sample.seed_tts_speech_extra + if not ex: + return + base = dict(rfi.extra_body) if rfi.extra_body else {} + base.update(ex) + rfi.extra_body = base + # Used by request funcs to force streaming TTS behavior and to export PCM when WER is on. + setattr(rfi, "seed_tts_row", True) + sys_prompt = (sample.seed_tts_system_prompt or "").strip() or SEED_TTS_DEFAULT_OMNI_SYSTEM_PROMPT + setattr( + rfi, + "omni_chat_messages", + [ + {"role": "system", "content": [{"type": "text", "text": sys_prompt}]}, + {"role": "user", "content": [{"type": "text", "text": sample.prompt}]}, + ], + ) + + +def _daily_omni_repo_from_args(args) -> str | None: + """Resolve HuggingFace repo id for Daily-Omni from CLI args. + + vLLM allows ``--dataset-path`` to be a local path while the real HF id is + passed via ``--hf-name``. Upstream ``get_samples`` for ``hf`` only matches + a fixed elif-chain and never discovers Omni's loader, so we must detect + Daily-Omni here using either field. + """ + dp = getattr(args, "dataset_path", None) + hn = getattr(args, "hf_name", None) + if dp in DailyOmniDataset.SUPPORTED_DATASET_PATHS: + return dp + if hn in DailyOmniDataset.SUPPORTED_DATASET_PATHS: + return hn + return None + def get_samples(args, tokenizer): - if args.backend not in ["openai-chat-omni", "openai-audio-speech"]: + # Daily-Omni: explicit dataset name, or hf + matching path/hf-name + is_daily_omni = args.dataset_name == "daily-omni" or ( + args.dataset_name == "hf" and _daily_omni_repo_from_args(args) is not None + ) + is_seed_tts = args.dataset_name == "seed-tts" + + # Check if we need to handle omni-related backends/datasets + is_omni_backend = args.backend in ["openai-chat-omni", "openai-audio-speech", "daily-omni"] + is_omni_dataset = is_daily_omni or is_seed_tts or args.dataset_name == "random-mm" + + if not is_omni_backend and not is_omni_dataset: + # Not an omni-related request, delegate to original implementation return get_samples_old(args, tokenizer) - elif args.dataset_name == "random-mm": + + # Handle Daily-Omni dataset + if is_daily_omni: + # Support: + # --dataset-name daily-omni [--dataset-path liarliar/Daily-Omni] + # --dataset-name daily-omni --daily-omni-qa-json /path/to/qa.json (offline QA) + # --dataset-name hf --dataset-path liarliar/Daily-Omni + # --dataset-name hf --hf-name liarliar/Daily-Omni (dataset-path may be local) + + # Validate backend supports multimodal (video) + if args.backend not in ["openai-chat-omni", "daily-omni"]: + raise ValueError( + f"Daily-Omni dataset requires a multimodal backend that supports video. " + f"Got backend='{args.backend}'. Please use '--backend openai-chat-omni'" + ) + + # Determine video directory if specified (for local video files) + video_dir = getattr(args, "daily_omni_video_dir", None) + + # Get HF split (default to "train"; unused when loading from local qa.json) + dataset_split = getattr(args, "hf_split", None) or "train" + + qa_json = getattr(args, "daily_omni_qa_json", None) + if isinstance(qa_json, str): + qa_json = qa_json.strip() or None + + if qa_json is not None: + logger.info( + "Loading Daily-Omni dataset: qa_json=%s, video_dir=%s (Hub not used for QA)", + qa_json, + video_dir, + ) + dataset = DailyOmniDataset( + qa_json_path=qa_json, + dataset_path=None, + dataset_split=dataset_split, + random_seed=args.seed, + video_dir=video_dir, + input_mode=getattr(args, "daily_omni_input_mode", "all"), + inline_local_video=getattr(args, "daily_omni_inline_local_video", False), + trust_remote_code=getattr(args, "trust_remote_code", False), + disable_shuffle=getattr(args, "disable_shuffle", False), + ) + else: + repo_id = _daily_omni_repo_from_args(args) + if args.dataset_name == "daily-omni": + if repo_id is None: + repo_id = _DEFAULT_DAILY_OMNI_REPO + elif repo_id is None: + raise ValueError( + "Daily-Omni with --dataset-name hf requires " + f"--dataset-path {_DEFAULT_DAILY_OMNI_REPO} or " + f"--hf-name {_DEFAULT_DAILY_OMNI_REPO}." + ) + + logger.info( + "Loading Daily-Omni dataset: hf_repo=%s, split=%s, video_dir=%s", + repo_id, + dataset_split, + video_dir, + ) + + dataset = DailyOmniDataset( + dataset_path=repo_id, + dataset_split=dataset_split, + dataset_subset=getattr(args, "hf_subset", None), + random_seed=args.seed, + video_dir=video_dir, + input_mode=getattr(args, "daily_omni_input_mode", "all"), + inline_local_video=getattr(args, "daily_omni_inline_local_video", False), + trust_remote_code=getattr(args, "trust_remote_code", False), + no_stream=getattr(args, "no_stream", False), + disable_shuffle=getattr(args, "disable_shuffle", False), + ) + + out_len = getattr(args, "output_len", None) + if out_len is None: + out_len = getattr(args, "hf_output_len", None) + if out_len is None: + out_len = DailyOmniDataset.DEFAULT_OUTPUT_LEN + + input_requests = dataset.sample( + tokenizer=tokenizer, + num_requests=args.num_prompts, + output_len=out_len, + request_id_prefix=args.request_id_prefix, + no_oversample=args.no_oversample, + ) + return input_requests + + if is_seed_tts: + if args.backend not in ("openai-audio-speech", "openai-chat-omni"): + raise ValueError( + "Seed-TTS requires --backend openai-audio-speech (POST /v1/audio/speech) or " + "--backend openai-chat-omni (POST /v1/chat/completions with ref_audio/ref_text). " + f"Got backend={args.backend!r}." + ) + repo_id = getattr(args, "dataset_path", None) or getattr(args, "hf_name", None) + if not repo_id: + raise ValueError( + "Seed-TTS requires --dataset-path (HF dataset repo id or local directory) or " + "--hf-name for the Hub dataset id." + ) + + dataset = SeedTTSDataset( + dataset_path=repo_id, + random_seed=args.seed, + locale=getattr(args, "seed_tts_locale", "en"), + inline_ref_audio=not getattr(args, "seed_tts_file_ref_audio", False), + seed_tts_root=getattr(args, "seed_tts_root", None), + system_prompt=getattr(args, "seed_tts_system_prompt", None), + disable_shuffle=getattr(args, "disable_shuffle", False), + ) + out_len = getattr(args, "output_len", None) + if out_len is None: + out_len = getattr(args, "hf_output_len", None) + if out_len is None: + out_len = SeedTTSDataset.DEFAULT_OUTPUT_LEN + return dataset.sample( + tokenizer=tokenizer, + num_requests=args.num_prompts, + output_len=out_len, + request_id_prefix=args.request_id_prefix, + no_oversample=args.no_oversample, + ) + + # Handle random-mm dataset (Omni's synthetic multimodal dataset) + if args.dataset_name == "random-mm": dataset = OmniRandomMultiModalDataset(random_seed=args.seed, dataset_path=args.dataset_path) input_requests = dataset.sample( tokenizer=tokenizer, @@ -64,6 +295,10 @@ def get_samples(args, tokenizer): datasets.get_samples = get_samples +_serve_mod = sys.modules.get("vllm.benchmarks.serve") +if _serve_mod is not None: + _serve_mod.get_samples = get_samples + @dataclass class MixRequestFuncOutput(RequestFuncOutput): @@ -72,6 +307,9 @@ class MixRequestFuncOutput(RequestFuncOutput): audio_frames: int = 0 audio_rtf: float = 0.0 text_latency: float = 0.0 + #: Raw PCM s16le mono at 24 kHz for Seed-TTS WER: from ``/v1/audio/speech`` stream or + #: resampled export after ``openai-chat-omni`` audio deltas. + tts_output_pcm_bytes: bytes | None = None async def async_request_openai_chat_omni_completions( @@ -83,13 +321,17 @@ async def async_request_openai_chat_omni_completions( api_url = request_func_input.api_url _validate_api_url(api_url, "OpenAI Chat Completions API", "chat/completions") - content = _get_chat_content(request_func_input, mm_position=mm_position) + omni_messages = getattr(request_func_input, "omni_chat_messages", None) + if omni_messages is not None: + messages_payload = omni_messages + else: + effective_mm_position = getattr(request_func_input, "mm_position", mm_position) + content = _get_chat_content(request_func_input, mm_position=effective_mm_position) + messages_payload = [{"role": "user", "content": content}] payload = { "model": request_func_input.model_name if request_func_input.model_name else request_func_input.model, - "messages": [ - {"role": "user", "content": content}, - ], + "messages": messages_payload, "temperature": 0.0, "max_tokens": request_func_input.output_len, "stream": True, @@ -98,6 +340,10 @@ async def async_request_openai_chat_omni_completions( }, } _update_payload_common(payload, request_func_input) + # Seed-TTS via chat: voice-clone fields live on the body; ensure audio is streamed. + if getattr(request_func_input, "seed_tts_row", False): + if payload.get("modalities") is None: + payload["modalities"] = ["text", "audio"] response_format = payload.get("response_format", "wav") if response_format == "pcm": @@ -167,7 +413,10 @@ async def async_request_openai_chat_omni_completions( data = json.loads(chunk) if choices := data.get("choices"): modality = data.get("modality") - content = choices[0]["delta"].get("content") + delta = choices[0].get("delta") or {} + content = delta.get("content") + if not content and isinstance(delta.get("audio"), dict): + content = delta["audio"].get("data") if modality == "text": # First token if ttft == 0.0: @@ -182,7 +431,7 @@ async def async_request_openai_chat_omni_completions( if output.audio_ttfp == 0.0: output.audio_ttfp = timestamp - st audio_generate_time = timestamp - st - if content != "": + if content: audio_bytes = base64.b64decode(content) seg = AudioSegment.from_file(io.BytesIO(audio_bytes)) if seg is not None: @@ -214,6 +463,12 @@ async def async_request_openai_chat_omni_completions( else: output.audio_rtf = 0 logger.warning("Audio duration is zero") + if _seed_tts_capture_pcm_for_wer() and getattr(request_func_input, "seed_tts_row", False): + try: + seg = generated_audio.set_frame_rate(24000).set_channels(1).set_sample_width(2) + output.tts_output_pcm_bytes = bytes(seg.raw_data) + except Exception as ex: + logger.warning("seed_tts WER PCM export failed: %s", ex) output.success = True else: output.error = response.reason or "" @@ -268,6 +523,10 @@ async def async_request_openai_audio_speech( "response_format": "pcm", } _update_payload_common(payload, request_func_input) + # Seed-TTS + WER: ``--extra-body`` may set stream=false / other formats; speech must stream PCM. + if getattr(request_func_input, "seed_tts_row", False) and _seed_tts_capture_pcm_for_wer(): + payload["stream"] = True + payload["response_format"] = "pcm" headers = { "Content-Type": "application/json", @@ -286,6 +545,8 @@ async def async_request_openai_audio_speech( st = time.perf_counter() output.start_time = st total_pcm_bytes = 0 + capture_wer_pcm = _seed_tts_capture_pcm_for_wer() and getattr(request_func_input, "seed_tts_row", False) + pcm_capture = bytearray() if capture_wer_pcm else None try: async with session.post(url=api_url, json=payload, headers=headers) as response: if response.status == 200: @@ -297,6 +558,8 @@ async def async_request_openai_audio_speech( output.audio_ttfp = timestamp - st output.ttft = output.audio_ttfp total_pcm_bytes += len(chunk) + if pcm_capture is not None: + pcm_capture.extend(chunk) end_time = time.perf_counter() output.latency = end_time - st @@ -309,6 +572,16 @@ async def async_request_openai_audio_speech( else: output.audio_rtf = 0 logger.warning("Audio duration is zero") + if pcm_capture is not None and pcm_capture: + output.tts_output_pcm_bytes = bytes(pcm_capture) + elif capture_wer_pcm: + ct = response.headers.get("Content-Type", "") + logger.warning( + "Seed-TTS WER: HTTP 200 but no PCM bytes (Content-Type=%r, url=%s). " + "Check stream=true and response_format=pcm on the server.", + ct, + api_url, + ) output.success = True else: output.error = response.reason or "" @@ -331,6 +604,12 @@ ASYNC_REQUEST_FUNCS["openai-audio-speech"] = async_request_openai_audio_speech if "openai-audio-speech" not in OPENAI_COMPATIBLE_BACKENDS: OPENAI_COMPATIBLE_BACKENDS.append("openai-audio-speech") +# Daily-Omni backend for audio-visual reasoning benchmark +# Reuses openai-chat-omni completions for video+text understanding +ASYNC_REQUEST_FUNCS["daily-omni"] = async_request_openai_chat_omni_completions +if "daily-omni" not in OPENAI_COMPATIBLE_BACKENDS: + OPENAI_COMPATIBLE_BACKENDS.append("daily-omni") + # ruff: noqa: E402 # Prevent import order from causing patch failures from vllm.benchmarks import serve @@ -422,6 +701,8 @@ async def benchmark( extra_headers=extra_headers, extra_body=extra_body, ) + _attach_daily_omni_to_request_func_input(input_requests[0], test_input) + _attach_seed_tts_to_request_func_input(input_requests[0], test_input) if ready_check_timeout_sec > 0: test_output = await wait_for_endpoint( @@ -484,6 +765,8 @@ async def benchmark( extra_headers=extra_headers, extra_body=extra_body, ) + _attach_daily_omni_to_request_func_input(input_requests[0], profile_input) + _attach_seed_tts_to_request_func_input(input_requests[0], profile_input) profile_output = await request_func(request_func_input=profile_input, session=session) if profile_output.success: print("Profiler started") @@ -564,6 +847,8 @@ async def benchmark( extra_body=extra_body, request_id=request_id, ) + _attach_daily_omni_to_request_func_input(request, request_func_input) + _attach_seed_tts_to_request_func_input(request, request_func_input) tasks.append( asyncio.create_task(limited_request_func(request_func_input=request_func_input, session=session, pbar=pbar)) ) @@ -631,6 +916,37 @@ async def benchmark( "errors": [output.error for output in outputs], } + from vllm_omni.benchmarks.data_modules.daily_omni_eval import ( + compute_daily_omni_accuracy_metrics, + print_daily_omni_accuracy_summary, + ) + + _save_items = os.environ.get("DAILY_OMNI_SAVE_EVAL_ITEMS", "").lower() in ( + "1", + "true", + "yes", + ) + _daily_acc = compute_daily_omni_accuracy_metrics(input_requests, outputs, include_per_item=_save_items) + if _daily_acc is not None: + result.update(_daily_acc) + print_daily_omni_accuracy_summary(_daily_acc) + + if _seed_tts_capture_pcm_for_wer(): + from vllm_omni.benchmarks.data_modules.seed_tts_eval import ( + compute_seed_tts_wer_metrics, + print_seed_tts_wer_summary, + ) + + _save_wer = os.environ.get("SEED_TTS_WER_SAVE_ITEMS", "").lower() in ( + "1", + "true", + "yes", + ) + _wer_m = compute_seed_tts_wer_metrics(input_requests, outputs, include_per_item=_save_wer) + if _wer_m is not None: + result.update(_wer_m) + print_seed_tts_wer_summary(_wer_m) + if rps_change_events: result["rps_change_events"] = rps_change_events diff --git a/vllm_omni/benchmarks/serve.py b/vllm_omni/benchmarks/serve.py index fe9460369..d3f3510c5 100644 --- a/vllm_omni/benchmarks/serve.py +++ b/vllm_omni/benchmarks/serve.py @@ -1,9 +1,21 @@ import argparse import asyncio +import os from typing import Any from vllm.benchmarks.serve import main_async +# Import patch to register daily-omni dataset and omni backends +# This monkey-patches vllm.benchmarks.datasets.get_samples before it's used +# Must be imported before any vllm.benchmarks module usage +import vllm_omni.benchmarks.patch.patch # noqa: F401 + def main(args: argparse.Namespace) -> dict[str, Any]: + if getattr(args, "seed_tts_wer_eval", False): + os.environ["SEED_TTS_WER_EVAL"] = "1" + if getattr(args, "seed_tts_wer_save_items", False): + os.environ["SEED_TTS_WER_SAVE_ITEMS"] = "1" + if getattr(args, "daily_omni_save_eval_items", False): + os.environ["DAILY_OMNI_SAVE_EVAL_ITEMS"] = "1" return asyncio.run(main_async(args)) diff --git a/vllm_omni/entrypoints/cli/benchmark/serve.py b/vllm_omni/entrypoints/cli/benchmark/serve.py index 906e8851a..d281432e5 100644 --- a/vllm_omni/entrypoints/cli/benchmark/serve.py +++ b/vllm_omni/entrypoints/cli/benchmark/serve.py @@ -1,4 +1,5 @@ import argparse +import os from vllm.benchmarks.serve import add_cli_args @@ -6,15 +7,149 @@ from vllm_omni.benchmarks.serve import main from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase +def add_daily_omni_cli_args(parser: argparse.ArgumentParser) -> None: + """Add CLI arguments specific to Daily-Omni dataset. + + This function should be called by the CLI entrypoint to add additional + arguments for daily-omni benchmark support. + + Args: + parser: The ArgumentParser instance to extend + """ + # Daily-Omni specific arguments + daily_omni_group = parser.add_argument_group("Daily-Omni Dataset Options") + + daily_omni_group.add_argument( + "--daily-omni-qa-json", + type=str, + default=None, + help="Path to local upstream qa.json. When set, QA rows are read from this file and " + "the HuggingFace dataset is not loaded (no network). Use with --daily-omni-video-dir " + "for fully offline runs. --dataset-path / Hub split flags are then ignored for QA loading.", + ) + daily_omni_group.add_argument( + "--daily-omni-video-dir", + type=str, + default=None, + help="Root directory of extracted Daily-Omni videos (contents of Videos.tar: " + "each video_id in its own subdir with {video_id}_video.mp4). " + "When using file URLs, you MUST start the vLLM server with " + "--allowed-local-media-path set to this same directory (or a parent), " + "otherwise requests fail with 'Cannot load local files without " + "--allowed-local-media-path'.", + ) + daily_omni_group.add_argument( + "--daily-omni-inline-local-video", + action="store_true", + default=False, + help="For local videos only: embed MP4 as base64 data URLs in benchmark " + "requests so the server does not need --allowed-local-media-path. " + "Increases request size and client memory; use for small --num-prompts. " + "When using --daily-omni-input-mode audio or all, local WAV files are " + "embedded the same way.", + ) + daily_omni_group.add_argument( + "--daily-omni-input-mode", + type=str, + choices=["all", "visual", "audio"], + default="all", + help="Daily-Omni input protocol (mirrors upstream Lliar-liar/Daily-Omni " + "--input_mode). 'visual': video only (default). 'audio': WAV only, " + "requires {video_id}/{video_id}_audio.wav under --daily-omni-video-dir. " + "'all': video + WAV together. Sets mm_processor_kwargs.use_audio_in_video=false " + "and matches official separate video/audio streams.", + ) + daily_omni_group.add_argument( + "--daily-omni-save-eval-items", + action="store_true", + default=False, + help="Include per-request Daily-Omni accuracy rows (gold/predicted/correct) " + "in the saved JSON under key daily_omni_eval_items. " + "Alternatively set env DAILY_OMNI_SAVE_EVAL_ITEMS=1.", + ) + + # Note: --dataset-name daily-omni via get_samples patch; use either Hub (--dataset-path + # liarliar/Daily-Omni) or local --daily-omni-qa-json (offline). + + +def add_seed_tts_cli_args(parser: argparse.ArgumentParser) -> None: + """CLI for Seed-TTS zero-shot TTS benchmark (``--dataset-name seed-tts``).""" + g = parser.add_argument_group("Seed-TTS Dataset Options") + g.add_argument( + "--seed-tts-locale", + type=str, + choices=["en", "zh"], + default="en", + help="Which Seed-TTS split to load: en/meta.lst or zh/meta.lst under the dataset root.", + ) + g.add_argument( + "--seed-tts-root", + type=str, + default=None, + help="Override root directory that contains en/ and zh/ (meta.lst + prompt-wavs). " + "If set, --dataset-path can still name the HF repo for logging; this path is used for files.", + ) + g.add_argument( + "--seed-tts-file-ref-audio", + action="store_true", + default=False, + help="Send ref_audio as file:// URIs (smaller HTTP bodies). Requires the API server " + "to be started with --allowed-local-media-path covering the Seed-TTS dataset root. " + "Default is inline data:audio/wav;base64 so Qwen3-TTS works without that flag.", + ) + g.add_argument( + "--seed-tts-inline-ref-audio", + action="store_true", + default=False, + help=argparse.SUPPRESS, + ) + g.add_argument( + "--seed-tts-system-prompt", + type=str, + default=None, + help="Override chat system message for --backend openai-chat-omni (Qwen3-Omni TTS). " + "Default follows official Qwen3-Omni identity + zero-shot voice-clone instructions.", + ) + g.add_argument( + "--seed-tts-wer-eval", + action="store_true", + default=False, + help="Keep synthesized audio as 24 kHz mono PCM for WER (works with " + "--backend openai-audio-speech or openai-chat-omni). Scoring follows " + "BytedanceSpeech/seed-tts-eval (Whisper-large-v3 / Paraformer-zh + jiwer). " + "Sets SEED_TTS_WER_EVAL=1. Install: pip install 'vllm-omni[seed-tts-eval]'. " + "Optional: SEED_TTS_EVAL_DEVICE, SEED_TTS_HF_WHISPER_MODEL.", + ) + g.add_argument( + "--seed-tts-wer-save-items", + action="store_true", + default=False, + help="Include per-utterance ASR rows in the saved JSON under key seed_tts_wer_eval_items. " + "Or set SEED_TTS_WER_SAVE_ITEMS=1.", + ) + + class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase): """The `serve` subcommand for vllm bench.""" name = "serve" - help = "Benchmark the online serving throughput." + help = "Benchmark the online serving throughput. Supports Daily-Omni and Seed-TTS datasets." @classmethod def add_cli_args(cls, parser: argparse.ArgumentParser) -> None: add_cli_args(parser) + + # Add Daily-Omni specific arguments + add_daily_omni_cli_args(parser) + add_seed_tts_cli_args(parser) + + for action in parser._actions: + if action.dest == "dataset_name" and action.choices is not None: + extra = [c for c in ("daily-omni", "seed-tts") if c not in action.choices] + if extra: + action.choices = list(action.choices) + extra + + # Update help messages for omni-specific features for action in parser._actions: if action.dest == "percentile_metrics": action.help = ( @@ -48,4 +183,10 @@ class OmniBenchmarkServingSubcommand(OmniBenchmarkSubcommandBase): @staticmethod def cmd(args: argparse.Namespace) -> None: + if getattr(args, "daily_omni_save_eval_items", False): + os.environ["DAILY_OMNI_SAVE_EVAL_ITEMS"] = "1" + if getattr(args, "seed_tts_wer_eval", False): + os.environ["SEED_TTS_WER_EVAL"] = "1" + if getattr(args, "seed_tts_wer_save_items", False): + os.environ["SEED_TTS_WER_SAVE_ITEMS"] = "1" main(args) From 0d020739a7d85e2b2ec2d30f26d0d741b4f4fb98 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zeyu=20Huang=20=7C=20=E9=BB=83=E6=BE=A4=E5=AE=87?= <11222265+fhfuih@users.noreply.github.com> Date: Wed, 15 Apr 2026 11:13:11 +0800 Subject: [PATCH 05/97] [CI] qwen image edit L4 accuracy test (#2761) --- .buildkite/test-nightly-diffusion.yml | 40 ++++ pyproject.toml | 1 + tests/conftest.py | 10 +- tests/e2e/accuracy/conftest.py | 25 +++ tests/e2e/accuracy/test_qwen_image_edit.py | 232 +++++++++++++++++++++ tests/e2e/accuracy/utils.py | 74 +++++++ 6 files changed, 377 insertions(+), 5 deletions(-) create mode 100644 tests/e2e/accuracy/test_qwen_image_edit.py create mode 100644 tests/e2e/accuracy/utils.py diff --git a/.buildkite/test-nightly-diffusion.yml b/.buildkite/test-nightly-diffusion.yml index a520ca435..b5ba8a117 100644 --- a/.buildkite/test-nightly-diffusion.yml +++ b/.buildkite/test-nightly-diffusion.yml @@ -375,3 +375,43 @@ steps: hostPath: path: /mnt/hf-cache type: DirectoryOrCreate + + - label: ":full_moon: Diffusion · Qwen-Image · Accuracy Test" + key: nightly-qwen-image-accuracy + timeout_in_minutes: 180 + if: *nightly_or_pr_label + commands: + - pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level advanced_model + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate diff --git a/pyproject.toml b/pyproject.toml index 753e0e398..9b034a7c8 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -55,6 +55,7 @@ dev = [ "pyttsx3>=2.99", "opencc>=1.2.0", "mistune>=3.2.0", # for example tests + "torchmetrics>=1.4.0", # for accuracy similarity metrics ] demo = [ diff --git a/tests/conftest.py b/tests/conftest.py index adb87cbd7..4ad4706fc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2397,7 +2397,7 @@ class OpenAIClientHandler: image_url = item.get("image_url", {}).get("url") else: image_url_obj = getattr(item, "image_url", None) - image_url = hasattr(image_url_obj, "url", None) if image_url_obj else None + image_url = getattr(image_url_obj, "url", None) if image_url_obj else None if image_url and image_url.startswith("data:image"): b64_data = image_url.split(",", 1)[1] img = decode_b64_image(b64_data) @@ -2703,7 +2703,7 @@ class OpenAIClientHandler: return responses - def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[OmniResponse]: + def send_diffusion_request(self, request_config: dict[str, Any], request_num: int = 1) -> list[DiffusionResponse]: """ Send OpenAI requests for diffusion models. @@ -2711,9 +2711,9 @@ class OpenAIClientHandler: request_config: Request configuration dictionary containing parameters like model, messages request_num: Number of requests to send concurrently, defaults to 1 (single request) Returns: - List[OmniResponse]: List of response objects + List[DiffusionResponse]: List of response objects """ - responses = [] + responses: list[DiffusionResponse] = [] stream = request_config.get("stream", False) modalities = request_config.get("modalities", omit) # Most diffusion models don't require modalities param extra_body = request_config.get("extra_body", None) @@ -2876,7 +2876,7 @@ class OpenAIClientHandler: return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}" -@pytest.fixture +@pytest.fixture(scope="module") def openai_client(request: pytest.FixtureRequest, run_level: str): """Create OpenAIClientHandler fixture to facilitate communication with OmniServer with encapsulated request sending, concurrent requests, response handling, and validation.""" diff --git a/tests/e2e/accuracy/conftest.py b/tests/e2e/accuracy/conftest.py index 062750b3c..3d614b8cd 100644 --- a/tests/e2e/accuracy/conftest.py +++ b/tests/e2e/accuracy/conftest.py @@ -5,10 +5,13 @@ import shutil import subprocess from contextlib import contextmanager from dataclasses import dataclass +from io import BytesIO from pathlib import Path import pytest +import requests import torch +from PIL import Image from tests.conftest import OmniServer, OmniServerParams @@ -183,6 +186,28 @@ def accuracy_artifact_root() -> Path: return root +@pytest.fixture(scope="session") +def qwen_bear_image(accuracy_artifact_root: Path) -> Image.Image: + """Download the Qwen bear image from the URL and save it to the accuracy artifact root.""" + QWEN_BEAR_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/qwen-bear.png" + response = requests.get(QWEN_BEAR_IMAGE_URL, timeout=60) + response.raise_for_status() + image = Image.open(BytesIO(response.content)).convert("RGB") + image.save(accuracy_artifact_root / "qwen_bear.png") + return image + + +@pytest.fixture(scope="session") +def rabbit_image(accuracy_artifact_root: Path) -> Image.Image: + """Download the rabbit image from the URL and save it to the accuracy artifact root.""" + RABBIT_IMAGE_URL = "https://vllm-public-assets.s3.us-west-2.amazonaws.com/omni-assets/rabbit.png" + response = requests.get(RABBIT_IMAGE_URL, timeout=60) + response.raise_for_status() + image = Image.open(BytesIO(response.content)).convert("RGB") + image.save(accuracy_artifact_root / "rabbit.png") + return image + + def reset_artifact_dir(path: Path) -> Path: if path.exists(): shutil.rmtree(path) diff --git a/tests/e2e/accuracy/test_qwen_image_edit.py b/tests/e2e/accuracy/test_qwen_image_edit.py new file mode 100644 index 000000000..9a9701034 --- /dev/null +++ b/tests/e2e/accuracy/test_qwen_image_edit.py @@ -0,0 +1,232 @@ +from __future__ import annotations + +import gc +from pathlib import Path + +import pytest +import requests +import torch +from diffusers import QwenImageEditPipeline, QwenImageEditPlusPipeline +from PIL import Image + +from benchmarks.accuracy.common import decode_base64_image, pil_to_png_bytes +from tests.conftest import ( + OmniServer, + _run_post_test_cleanup, + _run_pre_test_cleanup, +) +from tests.e2e.accuracy.utils import assert_similarity, model_output_dir +from tests.utils import hardware_test + +SINGLE_MODEL = "Qwen/Qwen-Image-Edit" +MULTIPLE_MODEL = "Qwen/Qwen-Image-Edit-2509" +WIDTH = 512 +HEIGHT = 512 +NUM_INFERENCE_STEPS = 20 +TRUE_CFG_SCALE = 4.0 +SEED = 42 +SSIM_THRESHOLD = 0.94 +PSNR_THRESHOLD = 28.0 + +PROMPT_SINGLE_IMAGE = "The input is a 2D cartoon bear mascot. Restyle it into a painterly oil artwork with warm colors while preserving the main structure." +PROMPT_MULTIPLE_IMAGE = "Put the cartoon bear mascot and the furry rabbit into one coherent scene with a painterly oil artwork style and consistent lighting." +NEGATIVE_PROMPT = "low quality, blurry, artifacts, distortion" +SERVER_ARGS = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"] + + +def _run_vllm_omni_image_edit( + *, + omni_server: OmniServer, + prompt: str, + input_images: list[Image.Image], + output_path: Path, +) -> Image.Image: + response = requests.post( + f"http://{omni_server.host}:{omni_server.port}/v1/images/edits", + data={ + "model": omni_server.model, + "prompt": prompt, + "size": f"{WIDTH}x{HEIGHT}", + "n": 1, + "response_format": "b64_json", + "negative_prompt": NEGATIVE_PROMPT, + "num_inference_steps": NUM_INFERENCE_STEPS, + "true_cfg_scale": TRUE_CFG_SCALE, + "seed": SEED, + }, + files=[ + ("image", (f"image_{index}.png", pil_to_png_bytes(image), "image/png")) + for index, image in enumerate(input_images) + ], + timeout=600, + ) + response.raise_for_status() + payload = response.json() + assert len(payload["data"]) == 1 + image = decode_base64_image(payload["data"][0]["b64_json"]) + image.load() + image.save(output_path) + return image + + +def _run_diffusers_image_edit( + *, + model: str, + pipeline_class: type[QwenImageEditPipeline] | type[QwenImageEditPlusPipeline], + prompt: str, + input_images: list[Image.Image], + output_path: Path, +) -> Image.Image: + _run_pre_test_cleanup(enable_force=True) + pipe: QwenImageEditPipeline | QwenImageEditPlusPipeline | None = None + device = torch.device("cuda:0") + torch.cuda.set_device(device) + try: + images = input_images[0] if len(input_images) == 1 else input_images + pipe = pipeline_class.from_pretrained( + model, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + ).to(device) + pipe.set_progress_bar_config(disable=False) + generator = torch.Generator(device=device).manual_seed(SEED) + result = pipe( # pyright: ignore[reportCallIssue] + prompt=prompt, + image=images, + negative_prompt=NEGATIVE_PROMPT, + num_inference_steps=NUM_INFERENCE_STEPS, + true_cfg_scale=TRUE_CFG_SCALE, + width=WIDTH, + height=HEIGHT, + generator=generator, + ) + output_image = result.images[0].convert("RGB") # pyright: ignore[reportAttributeAccessIssue] + output_image.save(output_path) + return output_image + finally: + if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"): + pipe.maybe_free_model_hooks() + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + _run_post_test_cleanup(enable_force=True) + + +def _vllm_omni_output_single_image( + accuracy_artifact_root: Path, + qwen_bear_image: Image.Image, +) -> Image.Image: + output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL) + output_path = output_dir / "vllm_omni_single.png" + with OmniServer(model=SINGLE_MODEL, serve_args=SERVER_ARGS) as server: + output = _run_vllm_omni_image_edit( + omni_server=server, + prompt=PROMPT_SINGLE_IMAGE, + input_images=[qwen_bear_image], + output_path=output_path, + ) + return output + + +def _diffusers_output_single_image(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> Image.Image: + output_dir = model_output_dir(accuracy_artifact_root, SINGLE_MODEL) + output_path = output_dir / "diffusers_single.png" + return _run_diffusers_image_edit( + model=SINGLE_MODEL, + pipeline_class=QwenImageEditPipeline, + prompt=PROMPT_SINGLE_IMAGE, + input_images=[qwen_bear_image], + output_path=output_path, + ) + + +def _vllm_omni_output_multiple_image( + accuracy_artifact_root: Path, + qwen_bear_image: Image.Image, + rabbit_image: Image.Image, +) -> Image.Image: + output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL) + output_path = output_dir / "vllm_omni_multiple.png" + with OmniServer(model=MULTIPLE_MODEL, serve_args=SERVER_ARGS) as server: + output = _run_vllm_omni_image_edit( + omni_server=server, + prompt=PROMPT_MULTIPLE_IMAGE, + input_images=[qwen_bear_image, rabbit_image], + output_path=output_path, + ) + return output + + +def _diffusers_output_multiple_image( + accuracy_artifact_root: Path, qwen_bear_image: Image.Image, rabbit_image: Image.Image +) -> Image.Image: + output_dir = model_output_dir(accuracy_artifact_root, MULTIPLE_MODEL) + output_path = output_dir / "diffusers_multiple.png" + return _run_diffusers_image_edit( + model=MULTIPLE_MODEL, + pipeline_class=QwenImageEditPlusPipeline, + prompt=PROMPT_MULTIPLE_IMAGE, + input_images=[qwen_bear_image, rabbit_image], + output_path=output_path, + ) + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_qwen_image_edit_single_matches_diffusers( + accuracy_artifact_root: Path, + qwen_bear_image: Image.Image, +) -> None: + vllm_image = _vllm_omni_output_single_image( + accuracy_artifact_root=accuracy_artifact_root, + qwen_bear_image=qwen_bear_image, + ) + diffusers_image = _diffusers_output_single_image( + accuracy_artifact_root=accuracy_artifact_root, + qwen_bear_image=qwen_bear_image, + ) + assert_similarity( + model_name=SINGLE_MODEL, + vllm_image=vllm_image, + diffusers_image=diffusers_image, + width=WIDTH, + height=HEIGHT, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + ) + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +@pytest.mark.skip( + reason="Skipping as the second image seems to be ignored by the API. Will come back to this later after #2772 is merged." +) +def test_qwen_image_edit_multiple_matches_diffusers( + accuracy_artifact_root: Path, + qwen_bear_image: Image.Image, + rabbit_image: Image.Image, +) -> None: + vllm_image = _vllm_omni_output_multiple_image( + accuracy_artifact_root=accuracy_artifact_root, + qwen_bear_image=qwen_bear_image, + rabbit_image=rabbit_image, + ) + diffusers_image = _diffusers_output_multiple_image( + accuracy_artifact_root=accuracy_artifact_root, + qwen_bear_image=qwen_bear_image, + rabbit_image=rabbit_image, + ) + assert_similarity( + model_name=MULTIPLE_MODEL, + vllm_image=vllm_image, + diffusers_image=diffusers_image, + width=WIDTH, + height=HEIGHT, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + ) diff --git a/tests/e2e/accuracy/utils.py b/tests/e2e/accuracy/utils.py new file mode 100644 index 000000000..eb0eea757 --- /dev/null +++ b/tests/e2e/accuracy/utils.py @@ -0,0 +1,74 @@ +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest +import torch +from PIL import Image +from torchmetrics.image import PeakSignalNoiseRatio, StructuralSimilarityIndexMeasure + + +def model_output_dir(parent_dir: Path, model: str) -> Path: + safe_model_name = model.split("/")[-1].replace(".", "_") + path = parent_dir / safe_model_name + path.mkdir(parents=True, exist_ok=True) + return path + + +def assert_similarity( + *, + model_name: str, + vllm_image: Image.Image, + diffusers_image: Image.Image, + width: int, + height: int, + ssim_threshold: float, + psnr_threshold: float, +) -> None: + requested_size = (width, height) + if diffusers_image.size != requested_size: + pytest.skip( + "Skipping as diffusers baseline output is corrupt and not comparable: " + f"dimensions do not match requested size; requested={requested_size}, got={diffusers_image.size}." + ) + + assert vllm_image.size == diffusers_image.size, ( + f"Online and diffusers output sizes mismatch: online={vllm_image.size}, diffusers={diffusers_image.size}" + ) + + ssim_score, psnr_score = compute_image_ssim_psnr(prediction=vllm_image, reference=diffusers_image) + print(f"{model_name} similarity metrics:") + print(f" SSIM: value={ssim_score:.6f}, threshold>={ssim_threshold:.6f}, range=[-1, 1], higher_is_better=True") + print( + f" PSNR: value={psnr_score:.6f} dB, threshold>={psnr_threshold:.6f} dB, range=[0, +inf), higher_is_better=True" + ) + + assert ssim_score >= ssim_threshold, ( + f"SSIM below threshold for {model_name}: got {ssim_score:.6f}, expected >= {ssim_threshold:.6f}." + ) + assert psnr_score >= psnr_threshold, ( + f"PSNR below threshold for {model_name}: got {psnr_score:.6f}, expected >= {psnr_threshold:.6f}." + ) + + +def compute_image_ssim_psnr( + *, + prediction: Image.Image, + reference: Image.Image, +) -> tuple[float, float]: + pred_tensor = _pil_to_batched_tensor(prediction) + ref_tensor = _pil_to_batched_tensor(reference) + + ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0) + psnr_metric = PeakSignalNoiseRatio(data_range=1.0) + + ssim_value = float(ssim_metric(pred_tensor, ref_tensor).item()) + psnr_value = float(psnr_metric(pred_tensor, ref_tensor).item()) + return ssim_value, psnr_value + + +def _pil_to_batched_tensor(image: Image.Image) -> torch.Tensor: + array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 + tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) + return tensor From 61a3cbdff5785290501d711717e2b2e526ffe34f Mon Sep 17 00:00:00 2001 From: Samit <285365963@qq.com> Date: Wed, 15 Apr 2026 11:46:06 +0800 Subject: [PATCH 06/97] [Perf] Eliminate Hop 3 IPC overhead for single-stage diffusion via inline execution (#2736) Signed-off-by: samithuang <285365963@qq.com> Signed-off-by: Samit <285365963@qq.com> --- .../test_inline_stage_diffusion_client.py | 96 +++++ .../test_async_omni_engine_stage_init.py | 3 +- vllm_omni/diffusion/data.py | 43 +++ .../inline_stage_diffusion_client.py | 348 ++++++++++++++++++ vllm_omni/diffusion/stage_diffusion_client.py | 25 ++ vllm_omni/diffusion/stage_diffusion_proc.py | 46 +-- vllm_omni/engine/async_omni_engine.py | 2 + vllm_omni/engine/orchestrator.py | 17 + vllm_omni/engine/stage_init_utils.py | 8 +- .../entrypoints/openai/video_api_utils.py | 3 + vllm_omni/outputs.py | 3 + 11 files changed, 546 insertions(+), 48 deletions(-) create mode 100644 tests/diffusion/test_inline_stage_diffusion_client.py create mode 100644 vllm_omni/diffusion/inline_stage_diffusion_client.py diff --git a/tests/diffusion/test_inline_stage_diffusion_client.py b/tests/diffusion/test_inline_stage_diffusion_client.py new file mode 100644 index 000000000..385f39b12 --- /dev/null +++ b/tests/diffusion/test_inline_stage_diffusion_client.py @@ -0,0 +1,96 @@ +from __future__ import annotations + +import asyncio +from unittest.mock import MagicMock, patch + +import pytest + +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient +from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def mock_engine(): + with patch("vllm_omni.diffusion.inline_stage_diffusion_client.DiffusionEngine") as mock: + engine_instance = MagicMock() + mock.make_engine.return_value = engine_instance + yield engine_instance + + +@pytest.fixture +def client(mock_engine): + metadata = StageMetadata( + stage_id=0, + stage_type="diffusion", + engine_output_type="image", + is_comprehension=False, + requires_multimodal_data=False, + engine_input_source="prompt", + final_output=True, + final_output_type="image", + default_sampling_params={}, + custom_process_input_func=None, + model_stage=None, + runtime_cfg=None, + ) + with patch.object(InlineStageDiffusionClient, "_enrich_config"): + od_config = MagicMock(spec=OmniDiffusionConfig) + c = InlineStageDiffusionClient(model="test_model", od_config=od_config, metadata=metadata, batch_size=1) + yield c + c.shutdown() + + +@pytest.mark.asyncio +async def test_inline_dispatch_request_success(client, mock_engine): + # Setup mock engine step to return a successful result + mock_result = OmniRequestOutput.from_diffusion(request_id="req-1", images=[MagicMock()]) + mock_engine.step.return_value = [mock_result] + + sampling_params = OmniDiffusionSamplingParams() + await client.add_request_async("req-1", "A test prompt", sampling_params) + + # Wait for the task to be processed + for _ in range(10): + output = client.get_diffusion_output_nowait() + if output is not None: + break + await asyncio.sleep(0.01) + + assert output is not None + assert output.request_id == "req-1" + mock_engine.step.assert_called_once() + + +@pytest.mark.asyncio +async def test_inline_dispatch_request_error(client, mock_engine): + # Setup mock engine step to raise an exception + mock_engine.step.side_effect = RuntimeError("Engine failure") + + sampling_params = OmniDiffusionSamplingParams() + await client.add_request_async("req-err", "A test prompt", sampling_params) + + for _ in range(10): + output = client.get_diffusion_output_nowait() + if output is not None: + break + await asyncio.sleep(0.01) + + assert output is not None + assert output.request_id == "req-err" + assert output.error == "Engine failure" + assert not output.images + + +def test_inline_shutdown(client, mock_engine): + assert not client._shutting_down + + # Shutting down should cleanly cancel anything queued and close engine + client.shutdown() + + assert client._shutting_down + mock_engine.close.assert_called_once() diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 7b995fe70..84b0cb0be 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -100,10 +100,11 @@ def test_initialize_stages_passes_stage_init_timeout_to_diffusion_handshake(monk engine.log_stats = False engine.model = "dummy-model" engine.config_path = "dummy-config" - engine.num_stages = 1 + engine.num_stages = 2 engine.async_chunk = False engine.diffusion_batch_size = 1 engine.single_stage_mode = False + engine._omni_master_server = None engine.stage_configs = [types.SimpleNamespace(stage_id=0, stage_type="diffusion", engine_args={})] metadata = types.SimpleNamespace( diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index 56a891aa5..fca0a5bad 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -666,6 +666,49 @@ class OmniDiffusionConfig: def update_multimodal_support(self) -> None: self.supports_multimodal_inputs = self.model_class_name in {"QwenImageEditPlusPipeline"} + def enrich_config(self) -> None: + """Load model metadata from HuggingFace and populate config fields. + + Diffusers-style models expose ``model_index.json`` with ``_class_name``. + Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``, + so we fall back to reading that and mapping model_type manually. + """ + from vllm.transformers_utils.config import get_hf_file_to_dict + + try: + config_dict = get_hf_file_to_dict("model_index.json", self.model) + if config_dict is not None: + if self.model_class_name is None: + self.model_class_name = config_dict.get("_class_name", None) + self.update_multimodal_support() + + tf_config_dict = get_hf_file_to_dict("transformer/config.json", self.model) + self.tf_model_config = TransformerConfig.from_dict(tf_config_dict) + else: + raise FileNotFoundError("model_index.json not found") + except (AttributeError, OSError, ValueError, FileNotFoundError): + cfg = get_hf_file_to_dict("config.json", self.model) + if cfg is None: + raise ValueError(f"Could not find config.json or model_index.json for model {self.model}") + + self.tf_model_config = TransformerConfig.from_dict(cfg) + model_type = cfg.get("model_type") + architectures = cfg.get("architectures") or [] + + if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: + self.model_class_name = "BagelPipeline" + self.tf_model_config = TransformerConfig() + self.update_multimodal_support() + elif model_type == "nextstep": + if self.model_class_name is None: + self.model_class_name = "NextStep11Pipeline" + self.tf_model_config = TransformerConfig() + self.update_multimodal_support() + elif architectures and len(architectures) == 1: + self.model_class_name = architectures[0] + else: + raise + @classmethod def from_kwargs(cls, **kwargs: Any) -> "OmniDiffusionConfig": # Backwards-compatibility: older callers may use a diffusion-specific diff --git a/vllm_omni/diffusion/inline_stage_diffusion_client.py b/vllm_omni/diffusion/inline_stage_diffusion_client.py new file mode 100644 index 000000000..a33a3e956 --- /dev/null +++ b/vllm_omni/diffusion/inline_stage_diffusion_client.py @@ -0,0 +1,348 @@ +"""Inline Stage Diffusion Client for vLLM-Omni multi-stage runtime. + +Runs DiffusionEngine in a ThreadPoolExecutor inside the Orchestrator process +instead of spawning a separate StageDiffusionProc subprocess, eliminating ZMQ +IPC overhead. Used when there is only a single diffusion stage. +""" + +from __future__ import annotations + +import asyncio +import time +from concurrent.futures import ThreadPoolExecutor +from typing import TYPE_CHECKING, Any + +import torch +from PIL import Image +from vllm.logger import init_logger + +from vllm_omni.diffusion.data import DiffusionRequestAbortedError +from vllm_omni.diffusion.diffusion_engine import DiffusionEngine +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.engine.stage_init_utils import StageMetadata +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput + +if TYPE_CHECKING: + from vllm_omni.diffusion.data import OmniDiffusionConfig + from vllm_omni.inputs.data import OmniPromptType + +logger = init_logger(__name__) + + +class InlineStageDiffusionClient: + """Runs DiffusionEngine in a thread executor inside the Orchestrator.""" + + stage_type: str = "diffusion" + + def __init__( + self, + model: str, + od_config: OmniDiffusionConfig, + metadata: StageMetadata, + batch_size: int = 1, + ) -> None: + self.model = model + self.od_config = od_config + self.stage_id = metadata.stage_id + self.final_output = metadata.final_output + self.final_output_type = metadata.final_output_type + self.default_sampling_params = metadata.default_sampling_params + self.custom_process_input_func = metadata.custom_process_input_func + self.engine_input_source = metadata.engine_input_source + self.batch_size = batch_size + + self._enrich_config() + self._engine = DiffusionEngine.make_engine(self.od_config) + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="inline-diffusion") + + self._output_queue: asyncio.Queue[OmniRequestOutput] = asyncio.Queue() + self._tasks: dict[str, asyncio.Task] = {} + self._shutting_down = False + + logger.info( + "[InlineStageDiffusionClient] Stage-%s initialized inline (batch_size=%d)", + self.stage_id, + self.batch_size, + ) + + def _enrich_config(self) -> None: + """Load model metadata from HuggingFace and populate od_config fields.""" + self.od_config.enrich_config() + + # ------------------------------------------------------------------ + # Request processing + # ------------------------------------------------------------------ + + async def add_request_async( + self, + request_id: str, + prompt: OmniPromptType, + sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[int, dict[str, Any]] | None = None, + ) -> None: + task = asyncio.create_task( + self._dispatch_request( + request_id, + prompt, + sampling_params, + kv_sender_info, + ) + ) + self._tasks[request_id] = task + + async def _dispatch_request( + self, + request_id: str, + prompt: Any, + sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[str, Any] | None = None, + ) -> None: + try: + request = OmniDiffusionRequest( + prompts=[prompt], + sampling_params=sampling_params, + request_ids=[request_id], + request_id=request_id, + kv_sender_info=kv_sender_info, + ) + + loop = asyncio.get_running_loop() + results = await loop.run_in_executor(self._executor, self._engine.step, request) + result = results[0] + if not result.request_id: + result.request_id = request_id + + self._output_queue.put_nowait(result) + except DiffusionRequestAbortedError as e: + logger.info("request_id: %s aborted: %s", request_id, str(e)) + except Exception as e: + logger.exception("Diffusion request %s failed: %s", request_id, e) + error_output = OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + ) + error_output.error = str(e) + self._output_queue.put_nowait(error_output) + finally: + self._tasks.pop(request_id, None) + + async def add_batch_request_async( + self, + request_id: str, + prompts: list[OmniPromptType], + sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[int, dict[str, Any]] | None = None, + ) -> None: + task = asyncio.create_task( + self._dispatch_batch( + request_id, + prompts, + sampling_params, + kv_sender_info, + ) + ) + self._tasks[request_id] = task + + async def _dispatch_batch( + self, + request_id: str, + prompts: list[Any], + sampling_params: OmniDiffusionSamplingParams, + kv_sender_info: dict[str, Any] | None = None, + ) -> None: + try: + request = OmniDiffusionRequest( + prompts=prompts, + sampling_params=sampling_params, + request_ids=[f"{request_id}-{i}" for i in range(len(prompts))], + request_id=request_id, + kv_sender_info=kv_sender_info, + ) + + loop = asyncio.get_running_loop() + results = await loop.run_in_executor(self._executor, self._engine.step, request) + + all_images: list = [] + merged_mm: dict[str, Any] = {} + merged_metrics: dict[str, Any] = {} + merged_durations: dict[str, float] = {} + merged_custom: dict[str, Any] = {} + peak_mem = 0.0 + latents = None + trajectory_latents: list[torch.Tensor] | None = None + trajectory_timesteps: list[torch.Tensor] | None = None + trajectory_log_probs: torch.Tensor | None = None + trajectory_decoded: list[Image.Image] | None = None + final_output_type = "image" + + for r in results: + all_images.extend(r.images) + merged_mm.update(r._multimodal_output) + merged_metrics.update(r.metrics) + merged_durations.update(r.stage_durations) + merged_custom.update(r._custom_output) + peak_mem = max(peak_mem, r.peak_memory_mb) + if latents is None and r.latents is not None: + latents = r.latents + if trajectory_latents is None: + trajectory_latents = r.trajectory_latents + if trajectory_timesteps is None: + trajectory_timesteps = r.trajectory_timesteps + if trajectory_log_probs is None: + trajectory_log_probs = r.trajectory_log_probs + if trajectory_decoded is None: + trajectory_decoded = r.trajectory_decoded + if r.final_output_type != "image": + final_output_type = r.final_output_type + + result = OmniRequestOutput.from_diffusion( + request_id=request_id, + images=all_images, + prompt=prompts[0] if len(prompts) == 1 else None, + metrics=merged_metrics, + latents=latents, + trajectory_latents=trajectory_latents, + trajectory_timesteps=trajectory_timesteps, + trajectory_log_probs=trajectory_log_probs, + trajectory_decoded=trajectory_decoded, + custom_output=merged_custom or None, + multimodal_output=merged_mm or None, + final_output_type=final_output_type, + stage_durations=merged_durations, + peak_memory_mb=peak_mem, + ) + + self._output_queue.put_nowait(result) + except DiffusionRequestAbortedError as e: + logger.info("request_id: %s aborted: %s", request_id, str(e)) + except Exception as e: + logger.exception("Batch diffusion request %s failed: %s", request_id, e) + error_output = OmniRequestOutput.from_diffusion( + request_id=request_id, + images=[], + ) + error_output.error = str(e) + self._output_queue.put_nowait(error_output) + finally: + self._tasks.pop(request_id, None) + + def get_diffusion_output_nowait(self) -> OmniRequestOutput | None: + try: + return self._output_queue.get_nowait() + except asyncio.QueueEmpty: + return None + + async def abort_requests_async(self, request_ids: list[str]) -> None: + for rid in request_ids: + task = self._tasks.pop(rid, None) + if task: + task.cancel() + self._engine.abort(rid) + + async def collective_rpc_async( + self, + method: str, + timeout: float | None = None, + args: tuple[Any, ...] = (), + kwargs: dict[str, Any] | None = None, + ) -> Any: + loop = asyncio.get_running_loop() + + if method == "profile": + is_start = args[0] if args else True + profile_prefix = args[1] if len(args) > 1 else None + if is_start and profile_prefix is None: + profile_prefix = f"stage_{self.stage_id}_diffusion_{int(time.time())}" + return await loop.run_in_executor( + self._executor, + self._engine.profile, + is_start, + profile_prefix, + ) + + kwargs = kwargs or {} + + # LoRA methods + if method == "add_lora": + lora_request = args[0] if args else kwargs.get("lora_request") + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "add_lora", + timeout, + (), + {"lora_request": lora_request}, + None, + ) + return all(results) if isinstance(results, list) else results + + if method == "remove_lora": + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "remove_lora", + timeout, + args, + kwargs, + None, + ) + return all(results) if isinstance(results, list) else results + + if method == "list_loras": + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "list_loras", + timeout, + (), + {}, + None, + ) + if not isinstance(results, list): + return results or [] + merged: set[int] = set() + for part in results: + merged.update(part or []) + return sorted(merged) + + if method == "pin_lora": + lora_id = args[0] if args else kwargs.get("adapter_id") + results = await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + "pin_lora", + timeout, + (), + {"adapter_id": lora_id}, + None, + ) + return all(results) if isinstance(results, list) else results + + return await loop.run_in_executor( + self._executor, + self._engine.collective_rpc, + method, + timeout, + args, + kwargs, + None, + ) + + def shutdown(self) -> None: + self._shutting_down = True + + # Cancel all pending tasks + for task in self._tasks.values(): + task.cancel() + + try: + # Cancel queued futures and wait for the running one to complete deterministically + self._executor.shutdown(wait=True, cancel_futures=True) + except Exception: + pass + + try: + self._engine.close() + except Exception: + pass diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index 7e740dc89..480d113d1 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -34,6 +34,24 @@ if TYPE_CHECKING: logger = init_logger(__name__) +def create_diffusion_client( + model: str, + od_config: OmniDiffusionConfig, + metadata: StageMetadata, + stage_init_timeout: int, + batch_size: int = 1, + use_inline: bool = False, +) -> Any: + """Factory to create either an inline or out-of-process diffusion client.""" + if use_inline: + from vllm_omni.diffusion.inline_stage_diffusion_client import InlineStageDiffusionClient + + return InlineStageDiffusionClient(model, od_config, metadata, batch_size=batch_size) + return StageDiffusionClient( + model, od_config, metadata, stage_init_timeout=stage_init_timeout, batch_size=batch_size + ) + + class StageDiffusionClient: """Communicates with StageDiffusionProc via ZMQ for use inside the Orchestrator. @@ -154,6 +172,13 @@ class StageDiffusionClient: "error": True, "reason": error_msg, } + elif req_id is not None: + error_output = OmniRequestOutput.from_diffusion( + request_id=req_id, + images=[], + ) + error_output.error = error_msg + self._output_queue.put_nowait(error_output) # Fields that are subprocess-local and cannot be serialized across # process boundaries. They are recreated in the subprocess with diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index cef697630..eced444fd 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -19,12 +19,11 @@ import zmq import zmq.asyncio from PIL import Image from vllm.logger import init_logger -from vllm.transformers_utils.config import get_hf_file_to_dict from vllm.utils.network_utils import get_open_zmq_ipc_path, zmq_socket_ctx from vllm.utils.system_utils import get_mp_context from vllm.v1.utils import shutdown -from vllm_omni.diffusion.data import DiffusionRequestAbortedError, TransformerConfig +from vllm_omni.diffusion.data import DiffusionRequestAbortedError from vllm_omni.diffusion.diffusion_engine import DiffusionEngine from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.distributed.omni_connectors.utils.serialization import ( @@ -66,47 +65,8 @@ class StageDiffusionProc: logger.info("StageDiffusionProc initialized with model: %s", self._model) def _enrich_config(self) -> None: - """Load model metadata from HuggingFace and populate od_config fields. - - Diffusers-style models expose ``model_index.json`` with ``_class_name``. - Non-diffusers models (e.g. Bagel, NextStep) only have ``config.json``, - so we fall back to reading that and mapping model_type manually. - """ - od_config = self._od_config - - try: - config_dict = get_hf_file_to_dict("model_index.json", od_config.model) - if config_dict is not None: - if od_config.model_class_name is None: - od_config.model_class_name = config_dict.get("_class_name", None) - od_config.update_multimodal_support() - - tf_config_dict = get_hf_file_to_dict("transformer/config.json", od_config.model) - od_config.tf_model_config = TransformerConfig.from_dict(tf_config_dict) - else: - raise FileNotFoundError("model_index.json not found") - except (AttributeError, OSError, ValueError, FileNotFoundError): - cfg = get_hf_file_to_dict("config.json", od_config.model) - if cfg is None: - raise ValueError(f"Could not find config.json or model_index.json for model {od_config.model}") - - od_config.tf_model_config = TransformerConfig.from_dict(cfg) - model_type = cfg.get("model_type") - architectures = cfg.get("architectures") or [] - - if model_type == "bagel" or "BagelForConditionalGeneration" in architectures: - od_config.model_class_name = "BagelPipeline" - od_config.tf_model_config = TransformerConfig() - od_config.update_multimodal_support() - elif model_type == "nextstep": - if od_config.model_class_name is None: - od_config.model_class_name = "NextStep11Pipeline" - od_config.tf_model_config = TransformerConfig() - od_config.update_multimodal_support() - elif architectures and len(architectures) == 1: - od_config.model_class_name = architectures[0] - else: - raise + """Load model metadata from HuggingFace and populate od_config fields.""" + self._od_config.enrich_config() # ------------------------------------------------------------------ # Request processing diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 9609cf6e2..054d5342d 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -759,12 +759,14 @@ class AsyncOmniEngine: self._omni_master_server, ) else: + use_inline = True if self.num_stages == 1 else False stage_clients[stage_idx] = initialize_diffusion_stage( self.model, stage_cfg, metadata, stage_init_timeout=stage_init_timeout, batch_size=self.diffusion_batch_size, + use_inline=use_inline, ) logger.info( "[AsyncOmniEngine] Stage %s initialized (diffusion, batch_size=%d)", diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 386b545eb..0fdab9c0d 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -246,6 +246,23 @@ class Orchestrator: idle = False req_state = self.request_states.get(output.request_id) if req_state is not None: + if getattr(output, "error", None) is not None: + parent_id = self._companion_to_parent.get(output.request_id, output.request_id) + await self.output_async_queue.put( + { + "type": "error", + "request_id": parent_id, + "stage_id": stage_id, + "error": output.error, + } + ) + role_map = self._companion_map.get(parent_id, {}) + for cid in role_map.values(): + self.request_states.pop(cid, None) + self._cleanup_companion_state(parent_id) + self.request_states.pop(parent_id, None) + continue + stage_metrics = self._build_stage_metrics(stage_id, output.request_id, [output], req_state) await self._route_output(stage_id, output, req_state, stage_metrics) continue diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 158b4c547..bf40aa77c 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -530,6 +530,7 @@ def initialize_diffusion_stage( metadata: StageMetadata, stage_init_timeout: int, batch_size: int = 1, + use_inline: bool = False, ) -> Any: """Build a diffusion stage client. @@ -541,13 +542,12 @@ def initialize_diffusion_stage( batch_size: Maximum number of requests to batch together in the diffusion engine. Passed through to ``StageDiffusionClient`` and ultimately to ``AsyncOmni``. + use_inline: If True, uses the inline diffusion client instead of subprocess. """ - from vllm_omni.diffusion.stage_diffusion_client import StageDiffusionClient + from vllm_omni.diffusion.stage_diffusion_client import create_diffusion_client od_config = build_diffusion_config(model, stage_cfg, metadata) - return StageDiffusionClient( - model, od_config, metadata, stage_init_timeout=stage_init_timeout, batch_size=batch_size - ) + return create_diffusion_client(model, od_config, metadata, stage_init_timeout, batch_size, use_inline) def _shutdown_or_close_resource(resource: Any, resource_name: str, stage_id: int) -> None: diff --git a/vllm_omni/entrypoints/openai/video_api_utils.py b/vllm_omni/entrypoints/openai/video_api_utils.py index 193546979..3fb991225 100644 --- a/vllm_omni/entrypoints/openai/video_api_utils.py +++ b/vllm_omni/entrypoints/openai/video_api_utils.py @@ -227,6 +227,9 @@ def _encode_video_bytes( frames_np *= 255.0 frames_u8 = np.round(frames_np).astype(np.uint8) + # Ensure contiguous memory layout for faster PyAV muxing + frames_u8 = np.ascontiguousarray(frames_u8) + audio_np = _coerce_audio_to_numpy(audio) if audio is not None else None return mux_video_audio_bytes( diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index 2c2c1d21c..c02c0c142 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -100,6 +100,9 @@ class OmniRequestOutput: # memory usage info peak_memory_mb: float = 0.0 + # error handling + error: str | None = None + @classmethod def from_pipeline( cls, From 6c6551dff8856e8e936cf29b5886174e4b149e4a Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Wed, 15 Apr 2026 11:52:20 +0800 Subject: [PATCH 07/97] [Feature] feat: add video frame interpolation postprocess (#2555) Signed-off-by: David Chen <530634352@qq.com> --- docs/.nav.yml | 1 + .../diffusion/frame_interpolation.md | 92 ++++ docs/user_guide/diffusion_features.md | 9 +- .../examples/online_serving/image_to_video.md | 32 ++ .../examples/online_serving/text_to_video.md | 32 ++ .../openai_api/test_video_api_utils.py | 92 ++++ .../openai_api/test_video_server.py | 134 +++++- vllm_omni/diffusion/diffusion_engine.py | 23 +- .../models/wan2_2/pipeline_wan2_2.py | 16 +- .../models/wan2_2/pipeline_wan2_2_i2v.py | 16 +- .../models/wan2_2/pipeline_wan2_2_ti2v.py | 16 +- vllm_omni/diffusion/postprocess/__init__.py | 10 + .../postprocess/rife_interpolator.py | 440 ++++++++++++++++++ vllm_omni/entrypoints/openai/api_server.py | 8 + .../entrypoints/openai/protocol/videos.py | 23 + vllm_omni/entrypoints/openai/serving_video.py | 22 +- vllm_omni/inputs/data.py | 4 + 17 files changed, 961 insertions(+), 9 deletions(-) create mode 100644 docs/user_guide/diffusion/frame_interpolation.md create mode 100644 tests/entrypoints/openai_api/test_video_api_utils.py create mode 100644 vllm_omni/diffusion/postprocess/__init__.py create mode 100644 vllm_omni/diffusion/postprocess/rife_interpolator.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 86ce4a3b0..441ef9f52 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -64,6 +64,7 @@ nav: - FP8: user_guide/diffusion/quantization/fp8.md - Int8: user_guide/diffusion/quantization/int8.md - GGUF: user_guide/diffusion/quantization/gguf.md + - Frame Interpolation: user_guide/diffusion/frame_interpolation.md - Parallelism: - Overview: user_guide/diffusion/parallelism/overview.md - CFG Parallel: user_guide/diffusion/parallelism/cfg_parallel.md diff --git a/docs/user_guide/diffusion/frame_interpolation.md b/docs/user_guide/diffusion/frame_interpolation.md new file mode 100644 index 000000000..349af50c5 --- /dev/null +++ b/docs/user_guide/diffusion/frame_interpolation.md @@ -0,0 +1,92 @@ +# Frame Interpolation + +## Overview + +vLLM-Omni supports post-generation frame interpolation for supported video +diffusion pipelines. This feature inserts synthesized intermediate frames +between adjacent generated frames to improve temporal smoothness without +rerunning the diffusion denoising loop. + +Frame interpolation runs in the diffusion worker post-processing path instead +of the API server encoding path. This allows the interpolation step to reuse +the worker's current accelerator device and keeps the FastAPI event loop free +from heavy synchronous PyTorch work. + +For an input video with `N` generated frames and interpolation exponent `exp`, +the output frame count is: + +```text +(N - 1) * 2**exp + 1 +``` + +The output FPS is multiplied by `2**exp` so the clip duration remains close to +the original generated video. + +## Supported Pipelines + +Frame interpolation is currently supported for: + +- `WanPipeline` (Wan2.2 text-to-video) +- `WanImageToVideoPipeline` +- `Wan22TI2VPipeline` + +## Request Parameters + +The video APIs `/v1/videos` and `/v1/videos/sync` accept: + +| Parameter | Type | Default | Description | +|-----------|------|---------|-------------| +| `enable_frame_interpolation` | bool | `false` | Enable post-generation frame interpolation | +| `frame_interpolation_exp` | int | `1` | Interpolation exponent. `1=2x`, `2=4x`, etc. | +| `frame_interpolation_scale` | float | `1.0` | RIFE inference scale | +| `frame_interpolation_model_path` | str | `None` | Local directory or Hugging Face repo ID containing `flownet.pkl` | + +## Execution Flow + +For supported Wan2.2 pipelines, the execution order is: + +1. Diffusion worker finishes denoising and decodes the raw video tensor. +2. Worker-side model-specific post-processing runs. +3. If frame interpolation is enabled, RIFE interpolates the decoded video + tensor on the worker side and records a FPS multiplier in `custom_output`. +4. The API server receives the already-interpolated video and only performs + MP4 export. + +This design keeps interpolation close to the generated tensor and avoids +introducing another heavyweight GPU context in the API server process. + +## Example + +Start the server: + +```bash +vllm serve Wan-AI/Wan2.2-T2V-A14B-Diffusers --omni --port 8091 +``` + +Run a sync request with interpolation enabled: + +```bash +curl -X POST http://localhost:8091/v1/videos/sync \ + -F "prompt=A dog running through a park" \ + -F "num_frames=81" \ + -F "width=832" \ + -F "height=480" \ + -F "fps=16" \ + -F "num_inference_steps=40" \ + -F "guidance_scale=1.0" \ + -F "guidance_scale_2=1.0" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ + -F "seed=42" \ + -o sync_t2v_interpolated.mp4 +``` + +## Notes + +- This is a post-processing feature. It does not modify the diffusion denoising + schedule. +- Higher interpolation exponents increase post-processing time and memory usage. +- If the interpolation model weights are not available locally, + `frame_interpolation_model_path` may point to a Hugging Face repo containing + `flownet.pkl`. diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 31cd1500f..45953b852 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -14,7 +14,7 @@ vLLM-Omni supports various advanced features for diffusion models: - Acceleration: **cache methods**, **parallelism methods**, **startup optimizations** - Memory optimization: **cpu offloading**, **quantization** -- Extensions: **LoRA inference** +- Extensions: **LoRA inference**, **frame interpolation** - Execution modes: **step execution** ## Supported Features @@ -69,6 +69,7 @@ Extension methods add specialized capabilities to diffusion models beyond standa | Method | Description | Best For | |--------|-------------|----------| | **[LoRA Inference](diffusion/lora.md)** | Enables inference with Low-Rank Adaptation (LoRA) adapters weights | Reinforcement learning extensions | +| **[Frame Interpolation](diffusion/frame_interpolation.md)** | Inserts intermediate video frames after generation for smoother motion | Video generation pipelines that need higher temporal smoothness | ### Execution Modes @@ -143,6 +144,11 @@ The following tables show which models support each feature: | **HunyuanVideo-1.5 T2V I2V** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ✅ | ✅ (decode) | ✅ | ❌ | | **DreamID-Omni** | ❌ | ❌ | ❌ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | ❌ | +**Frame Interpolation Support** + +- **Supported**: Wan2.2 text-to-video, image-to-video, and TI2V pipelines +- **Not supported**: Wan2.1-VACE, LTX-2, Helios, HunyuanVideo-1.5, DreamID-Omni + ### AudioGen | Model | ⚡TeaCache | ⚡Cache-DiT | 🔀SP (Ulysses & Ring) | 🔀CFG-Parallel | 🔀Tensor-Parallel | 🔀HSDP | 💾CPU Offload (Layerwise) | 💾VAE-Patch-Parallel | 💾Quantization | 🔄Step Execution | @@ -258,6 +264,7 @@ Measured on NVIDIA H800: **Extensions:** - **[LoRA Inference Guide](diffusion/lora.md)** - Low-Rank Adaptation for style customization and fine-tuning +- **[Frame Interpolation Guide](diffusion/frame_interpolation.md)** - Worker-side post-generation video frame interpolation for smoother motion **Execution Modes:** diff --git a/docs/user_guide/examples/online_serving/image_to_video.md b/docs/user_guide/examples/online_serving/image_to_video.md index 00b67d74e..781f0c2a5 100644 --- a/docs/user_guide/examples/online_serving/image_to_video.md +++ b/docs/user_guide/examples/online_serving/image_to_video.md @@ -72,6 +72,9 @@ curl -X POST http://localhost:8091/v1/videos/sync \ -F "guidance_scale_2=1.0" \ -F "boundary_ratio=0.875" \ -F "flow_shift=12.0" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ -F "seed=42" \ -o sync_i2v_output.mp4 ``` @@ -114,6 +117,9 @@ create_response=$(curl -s http://localhost:8091/v1/videos \ -F "guidance_scale_2=1.0" \ -F "boundary_ratio=0.875" \ -F "flow_shift=12.0" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ -F "seed=42") video_id=$(echo "$create_response" | jq -r '.id') @@ -172,9 +178,35 @@ curl -X POST http://localhost:8091/v1/videos \ -F "guidance_scale_2=1.0" \ -F "boundary_ratio=0.875" \ -F "flow_shift=12.0" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ -F "seed=42" ``` +Frame interpolation is also available for supported Wan2.2 I2V requests. See +[Frame Interpolation](../../diffusion/frame_interpolation.md) for worker-side +execution details and feature constraints. + +### Frame Interpolation Example + +```bash +curl -X POST http://localhost:8091/v1/videos/sync \ + -F "prompt=A bear playing with yarn, smooth motion" \ + -F "input_reference=@/path/to/qwen-bear.png" \ + -F "width=832" \ + -F "height=480" \ + -F "num_frames=33" \ + -F "fps=16" \ + -F "num_inference_steps=40" \ + -F "guidance_scale=1.0" \ + -F "guidance_scale_2=1.0" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ + -o sync_i2v_interpolated.mp4 +``` + ## Create Response Format `POST /v1/videos` returns a job record, not inline base64 video data. diff --git a/docs/user_guide/examples/online_serving/text_to_video.md b/docs/user_guide/examples/online_serving/text_to_video.md index 01e6d9d46..00a9c1672 100644 --- a/docs/user_guide/examples/online_serving/text_to_video.md +++ b/docs/user_guide/examples/online_serving/text_to_video.md @@ -165,6 +165,9 @@ curl -X POST http://localhost:8091/v1/videos \ -F "guidance_scale_2=4.0" \ -F "boundary_ratio=0.875" \ -F "flow_shift=5.0" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ -F "seed=42" ``` @@ -187,6 +190,35 @@ curl -X POST http://localhost:8091/v1/videos \ | `flow_shift` | float | None | Scheduler flow shift (Wan2.2) | | `seed` | int | None | Random seed (reproducible) | | `lora` | object | None | LoRA configuration | +| `enable_frame_interpolation` | bool | false | Enable RIFE frame interpolation before MP4 encoding | +| `frame_interpolation_exp` | int | 1 | Interpolation exponent; 1=2x temporal resolution, 2=4x | +| `frame_interpolation_scale` | float | 1.0 | RIFE inference scale; use 0.5 for high-resolution inputs | +| `frame_interpolation_model_path` | str | None | Local directory or Hugging Face repo ID with `flownet.pkl`; defaults to `elfgum/RIFE-4.22.lite` | + +## Frame Interpolation + +Frame interpolation is an optional post-processing step for `/v1/videos` and +`/v1/videos/sync`. It synthesizes intermediate frames between generated frames +without rerunning the diffusion model. If the generated video has `N` frames, +the interpolated output frame count is `(N - 1) * 2**exp + 1`. The encoder FPS +is multiplied by `2**exp` so the output duration remains close to the original. + +Frame interpolation runs in the diffusion worker post-processing path instead of +the API server encoding path, so it can reuse the worker's current accelerator +device without blocking the FastAPI event loop. + +Example: generate 5 frames and interpolate to 9 frames: + +```bash +curl -X POST http://localhost:8091/v1/videos/sync \ + -F "prompt=A dog running through a park" \ + -F "num_frames=5" \ + -F "fps=8" \ + -F "enable_frame_interpolation=true" \ + -F "frame_interpolation_exp=1" \ + -F "frame_interpolation_scale=1.0" \ + -o sync_t2v_interpolated.mp4 +``` ## Create Response Format diff --git a/tests/entrypoints/openai_api/test_video_api_utils.py b/tests/entrypoints/openai_api/test_video_api_utils.py new file mode 100644 index 000000000..5012c9b98 --- /dev/null +++ b/tests/entrypoints/openai_api/test_video_api_utils.py @@ -0,0 +1,92 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for OpenAI-compatible video API encoding helpers.""" + +import numpy as np +import pytest +import torch + +from vllm_omni.diffusion.postprocess import rife_interpolator +from vllm_omni.entrypoints.openai import video_api_utils + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _install_fake_video_mux(monkeypatch, mux_calls): + def _fake_mux_video_audio_bytes(frames, audio, fps, audio_sample_rate, video_codec_options=None): + mux_calls.append( + { + "frames": frames, + "audio": audio, + "fps": fps, + "audio_sample_rate": audio_sample_rate, + "video_codec_options": video_codec_options, + } + ) + return b"fake-video" + + monkeypatch.setattr( + "vllm_omni.diffusion.utils.media_utils.mux_video_audio_bytes", + _fake_mux_video_audio_bytes, + ) + + +def test_encode_video_bytes_exports_frames_without_interpolation(monkeypatch): + mux_calls = [] + _install_fake_video_mux(monkeypatch, mux_calls) + + frames = [np.full((2, 2, 3), fill_value=i / 5, dtype=np.float32) for i in range(5)] + video_bytes = video_api_utils._encode_video_bytes( + frames, + fps=8, + ) + + assert video_bytes == b"fake-video" + assert mux_calls[0]["frames"].shape == (5, 2, 2, 3) + assert mux_calls[0]["frames"].dtype == np.uint8 + assert mux_calls[0]["fps"] == 8.0 + assert mux_calls[0]["audio"] is None + + +def test_rife_model_inference_runs_on_dummy_tensors(): + model = rife_interpolator.Model().eval() + img0 = torch.rand(1, 3, 32, 32) + img1 = torch.rand(1, 3, 32, 32) + + output = model.inference(img0, img1, scale=1.0) + + assert output.shape == (1, 3, 32, 32) + assert torch.isfinite(output).all() + + +def test_frame_interpolator_runs_actual_torch_tensor_path(monkeypatch): + model = rife_interpolator.Model().eval() + interpolator = rife_interpolator.FrameInterpolator() + monkeypatch.setattr(interpolator, "_ensure_model_loaded", lambda preferred_device=None: model) + + video = torch.zeros(1, 3, 2, 32, 32) + output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0) + + assert multiplier == 2 + assert output_video.shape == (1, 3, 3, 32, 32) + assert torch.isfinite(output_video).all() + + +def test_frame_interpolator_prefers_input_tensor_device(monkeypatch): + chosen_devices = [] + model = rife_interpolator.Model().eval() + + def _fake_ensure_model_loaded(*, preferred_device=None): + chosen_devices.append(preferred_device) + return model + + interpolator = rife_interpolator.FrameInterpolator() + monkeypatch.setattr(interpolator, "_ensure_model_loaded", _fake_ensure_model_loaded) + monkeypatch.setattr(model.flownet, "to", lambda device: model.flownet) + + video = torch.zeros(1, 3, 2, 32, 32) + output_video, multiplier = interpolator.interpolate_tensor(video, exp=1, scale=1.0) + + assert chosen_devices == [video.device] + assert multiplier == 2 + assert output_video.shape == (1, 3, 3, 32, 32) diff --git a/tests/entrypoints/openai_api/test_video_server.py b/tests/entrypoints/openai_api/test_video_server.py index 82c34f87e..7a395bab5 100644 --- a/tests/entrypoints/openai_api/test_video_server.py +++ b/tests/entrypoints/openai_api/test_video_server.py @@ -34,15 +34,28 @@ pytestmark = [pytest.mark.core_model, pytest.mark.cpu] class MockVideoResult: - def __init__(self, videos, audios=None, sample_rate=None, stage_durations=None, peak_memory_mb=0.0): + def __init__( + self, + videos, + audios=None, + sample_rate=None, + custom_output=None, + stage_durations=None, + peak_memory_mb=0.0, + ): self.multimodal_output = {"video": videos} if audios is not None: self.multimodal_output["audio"] = audios if sample_rate is not None: self.multimodal_output["audio_sample_rate"] = sample_rate + self._custom_output = custom_output or {} self.stage_durations = stage_durations or {} self.peak_memory_mb = peak_memory_mb + @property + def custom_output(self): + return self._custom_output + class FakeAsyncOmni: def __init__(self): @@ -400,6 +413,67 @@ def test_sampling_params_pass_through(test_client, mocker: MockerFixture): assert captured.extra_args["flow_shift"] == 0.25 +def test_frame_interpolation_params_pass_to_diffusion_sampling_params(test_client, mocker: MockerFixture): + """Frame interpolation parameters should be forwarded to diffusion worker sampling params.""" + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + return_value=b"fake-video", + ) + response = test_client.post( + "/v1/videos", + data={ + "prompt": "smooth motion", + "fps": "8", + "enable_frame_interpolation": "true", + "frame_interpolation_exp": "2", + "frame_interpolation_scale": "0.5", + "frame_interpolation_model_path": "local-rife", + }, + ) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + + engine = test_client.app.state.openai_serving_video._engine_client + captured = engine.captured_sampling_params_list[0] + assert captured.enable_frame_interpolation is True + assert captured.frame_interpolation_exp == 2 + assert captured.frame_interpolation_scale == 0.5 + assert captured.frame_interpolation_model_path == "local-rife" + + +def test_worker_fps_multiplier_is_applied_to_async_encoding(test_client, mocker: MockerFixture): + fps_values = [] + engine = test_client.app.state.openai_serving_video._engine_client + + async def _generate(prompt, request_id, sampling_params_list): + engine.captured_prompt = prompt + engine.captured_sampling_params_list = sampling_params_list + import numpy as np + + yield MockVideoResult([np.zeros((1, 64, 64, 3), dtype=np.uint8)], custom_output={"video_fps_multiplier": 2}) + + engine.generate = _generate + + def _fake_encode(video, fps, **kwargs): + del video, kwargs + fps_values.append(fps) + return b"fake-video" + + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + side_effect=_fake_encode, + ) + + response = test_client.post("/v1/videos", data={"prompt": "fps multiplier", "fps": "8"}) + + assert response.status_code == 200 + video_id = response.json()["id"] + _wait_for_status(test_client, video_id, VideoGenerationStatus.COMPLETED.value) + assert fps_values == [16] + + def test_audio_sample_rate_comes_from_model_config(test_client, mocker: MockerFixture): audio_sample_rates = [] @@ -595,6 +669,10 @@ def test_video_request_validation(): with pytest.raises(ValueError): VideoGenerationRequest(prompt="test", image_reference={"file_id": "file-1", "image_url": "https://example.com"}) + with pytest.raises(ValueError): + VideoGenerationRequest(prompt="test", frame_interpolation_exp=0) + with pytest.raises(ValueError): + VideoGenerationRequest(prompt="test", frame_interpolation_scale=0) def test_list_videos_supports_order_after_and_limit(test_client, mocker: MockerFixture): @@ -1032,3 +1110,57 @@ def test_sync_sampling_params_pass_through(test_client, mocker: MockerFixture): assert captured.num_inference_steps == 30 assert captured.guidance_scale == 6.5 assert captured.seed == 42 + + +def test_sync_frame_interpolation_params_pass_to_sampling_params(test_client, mocker: MockerFixture): + """Frame interpolation parameters should be forwarded on the sync path.""" + encode_mock = _mock_encode_video_bytes(mocker) + response = test_client.post( + "/v1/videos/sync", + data={ + "prompt": "smooth sync", + "fps": "8", + "enable_frame_interpolation": "true", + "frame_interpolation_exp": "2", + "frame_interpolation_scale": "0.5", + "frame_interpolation_model_path": "local-rife", + }, + ) + + assert response.status_code == 200 + engine = test_client.app.state.openai_serving_video._engine_client + captured = engine.captured_sampling_params_list[0] + assert captured.enable_frame_interpolation is True + assert captured.frame_interpolation_exp == 2 + assert captured.frame_interpolation_scale == 0.5 + assert captured.frame_interpolation_model_path == "local-rife" + _, kwargs = encode_mock.call_args + assert kwargs["fps"] == 8 + + +def test_worker_fps_multiplier_is_applied_to_sync_encoding(test_client, mocker: MockerFixture): + engine = test_client.app.state.openai_serving_video._engine_client + fps_values = [] + + async def _generate(prompt, request_id, sampling_params_list): + engine.captured_prompt = prompt + engine.captured_sampling_params_list = sampling_params_list + yield MockVideoResult([object()], custom_output={"video_fps_multiplier": 2}) + + engine.generate = _generate + + def _fake_encode(video, fps, **kwargs): + del video, kwargs + fps_values.append(fps) + return b"fps-multiplied" + + mocker.patch( + "vllm_omni.entrypoints.openai.serving_video._encode_video_bytes", + side_effect=_fake_encode, + ) + + response = test_client.post("/v1/videos/sync", data={"prompt": "fps multiplier", "fps": "8"}) + + assert response.status_code == 200 + assert response.content == b"fps-multiplied" + assert fps_values == [16] diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 52a8f3854..fe940d623 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -3,6 +3,7 @@ from __future__ import annotations +import inspect import queue import threading import time @@ -78,6 +79,12 @@ class DiffusionEngine: self.post_process_func = get_diffusion_post_process_func(od_config) self.pre_process_func = get_diffusion_pre_process_func(od_config) + # Cache whether the model-specific postprocess accepts request-level + # sampling params so step() can support both legacy and extended hooks. + self._post_process_accepts_sampling_params = bool( + self.post_process_func is not None + and "sampling_params" in inspect.signature(self.post_process_func).parameters + ) executor_class = DiffusionExecutor.get_class(od_config) self.executor = executor_class(od_config) @@ -143,12 +150,22 @@ class DiffusionEngine: output_data = output_data.cpu() postprocess_start_time = time.perf_counter() - outputs = self.post_process_func(output_data) if self.post_process_func is not None else output_data + if self.post_process_func is not None: + # Some video pipelines need request-level controls during + # postprocess (for example worker-side frame interpolation). + if self._post_process_accepts_sampling_params: + outputs = self.post_process_func(output_data, sampling_params=request.sampling_params) + else: + outputs = self.post_process_func(output_data) + else: + outputs = output_data audio_payload = None + custom_output = output.custom_output or {} model_audio_sample_rate = None model_fps = None if isinstance(outputs, dict): audio_payload = outputs.get("audio") + custom_output.update(outputs.get("custom_output") or {}) model_audio_sample_rate = outputs.get("audio_sample_rate") model_fps = outputs.get("fps") outputs = outputs.get("video", outputs) @@ -225,7 +242,7 @@ class DiffusionEngine: trajectory_timesteps=output.trajectory_timesteps, trajectory_log_probs=output.trajectory_log_probs, trajectory_decoded=output.trajectory_decoded, - custom_output=output.custom_output or {}, + custom_output=custom_output, multimodal_output=mm_output, stage_durations=output.stage_durations, peak_memory_mb=output.peak_memory_mb, @@ -295,7 +312,7 @@ class DiffusionEngine: trajectory_timesteps=output.trajectory_timesteps, trajectory_log_probs=output.trajectory_log_probs, trajectory_decoded=output.trajectory_decoded, - custom_output=output.custom_output or {}, + custom_output=custom_output, multimodal_output=mm_output, stage_durations=output.stage_durations, peak_memory_mb=output.peak_memory_mb, diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index 84d89619e..a1b10439c 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -26,6 +26,7 @@ from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_z from vllm_omni.diffusion.models.schedulers import FlowUniPCMultistepScheduler from vllm_omni.diffusion.models.wan2_2.scheduling_wan_euler import WanEulerScheduler from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3DModel +from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt @@ -162,10 +163,23 @@ def get_wan22_post_process_func( def post_process_func( video: torch.Tensor, output_type: str = "np", + sampling_params=None, ): if output_type == "latent": return video - return video_processor.postprocess_video(video, output_type=output_type) + custom_output = {} + if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): + video, multiplier = interpolate_video_tensor( + video, + exp=sampling_params.frame_interpolation_exp, + scale=sampling_params.frame_interpolation_scale, + model_path=sampling_params.frame_interpolation_model_path, + ) + custom_output["video_fps_multiplier"] = multiplier + return { + "video": video_processor.postprocess_video(video, output_type=output_type), + "custom_output": custom_output, + } return post_process_func diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index 46484cd78..ddc6e0bc2 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -32,6 +32,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( resolve_wan_sample_solver, retrieve_latents, ) +from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt @@ -74,10 +75,23 @@ def get_wan22_i2v_post_process_func( def post_process_func( video: torch.Tensor, output_type: str = "np", + sampling_params=None, ): if output_type == "latent": return video - return video_processor.postprocess_video(video, output_type=output_type) + custom_output = {} + if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): + video, multiplier = interpolate_video_tensor( + video, + exp=sampling_params.frame_interpolation_exp, + scale=sampling_params.frame_interpolation_scale, + model_path=sampling_params.frame_interpolation_model_path, + ) + custom_output["video_fps_multiplier"] = multiplier + return { + "video": video_processor.postprocess_video(video, output_type=output_type), + "custom_output": custom_output, + } return post_process_func diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py index 939fe294a..62df13cbd 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py @@ -44,6 +44,7 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( resolve_wan_sample_solver, retrieve_latents, ) +from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.request import OmniDiffusionRequest from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -61,10 +62,23 @@ def get_wan22_ti2v_post_process_func( def post_process_func( video: torch.Tensor, output_type: str = "np", + sampling_params=None, ): if output_type == "latent": return video - return video_processor.postprocess_video(video, output_type=output_type) + custom_output = {} + if sampling_params is not None and getattr(sampling_params, "enable_frame_interpolation", False): + video, multiplier = interpolate_video_tensor( + video, + exp=sampling_params.frame_interpolation_exp, + scale=sampling_params.frame_interpolation_scale, + model_path=sampling_params.frame_interpolation_model_path, + ) + custom_output["video_fps_multiplier"] = multiplier + return { + "video": video_processor.postprocess_video(video, output_type=output_type), + "custom_output": custom_output, + } return post_process_func diff --git a/vllm_omni/diffusion/postprocess/__init__.py b/vllm_omni/diffusion/postprocess/__init__.py new file mode 100644 index 000000000..e6fe5b2d2 --- /dev/null +++ b/vllm_omni/diffusion/postprocess/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Diffusion post-processing helpers.""" + +from vllm_omni.diffusion.postprocess.rife_interpolator import ( + FrameInterpolator, + interpolate_video_tensor, +) + +__all__ = ["FrameInterpolator", "interpolate_video_tensor"] diff --git a/vllm_omni/diffusion/postprocess/rife_interpolator.py b/vllm_omni/diffusion/postprocess/rife_interpolator.py new file mode 100644 index 000000000..b2b4a9319 --- /dev/null +++ b/vllm_omni/diffusion/postprocess/rife_interpolator.py @@ -0,0 +1,440 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +RIFE 4.22.lite frame interpolation for vLLM-Omni video generation. + +RIFE model code is vendored and adapted from: + - https://github.com/hzwer/ECCV2022-RIFE (MIT License) + - https://github.com/hzwer/Practical-RIFE (MIT License) + Copyright (c) 2021 Zhewei Huang + +The FrameInterpolator wrapper and vLLM-Omni integration code are original work. +""" + +from __future__ import annotations + +import os +import threading +from typing import Any + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger + +logger = init_logger(__name__) + +_DEFAULT_RIFE_HF_REPO = "elfgum/RIFE-4.22.lite" +_MODEL_CACHE: dict[tuple[str, str], Model] = {} +_MODEL_CACHE_LOCK = threading.Lock() + + +def warp(ten_input: torch.Tensor, ten_flow: torch.Tensor) -> torch.Tensor: + """Warp input tensor by optical flow using grid_sample.""" + ten_horizontal = ( + torch.linspace(-1.0, 1.0, ten_flow.shape[3], device=ten_flow.device) + .view(1, 1, 1, ten_flow.shape[3]) + .expand(ten_flow.shape[0], -1, ten_flow.shape[2], -1) + ) + ten_vertical = ( + torch.linspace(-1.0, 1.0, ten_flow.shape[2], device=ten_flow.device) + .view(1, 1, ten_flow.shape[2], 1) + .expand(ten_flow.shape[0], -1, -1, ten_flow.shape[3]) + ) + ten_grid = torch.cat([ten_horizontal, ten_vertical], dim=1) + + ten_flow = torch.cat( + [ + ten_flow[:, 0:1, :, :] / ((ten_input.shape[3] - 1.0) / 2.0), + ten_flow[:, 1:2, :, :] / ((ten_input.shape[2] - 1.0) / 2.0), + ], + dim=1, + ) + grid = (ten_grid + ten_flow).permute(0, 2, 3, 1) + return F.grid_sample( + input=ten_input, + grid=grid, + mode="bilinear", + padding_mode="border", + align_corners=True, + ) + + +def _conv( + in_planes: int, + out_planes: int, + kernel_size: int = 3, + stride: int = 1, + padding: int = 1, + dilation: int = 1, +) -> nn.Sequential: + return nn.Sequential( + nn.Conv2d( + in_planes, + out_planes, + kernel_size=kernel_size, + stride=stride, + padding=padding, + dilation=dilation, + bias=True, + ), + nn.LeakyReLU(0.2, True), + ) + + +class ResConv(nn.Module): + """Residual convolution block with learnable beta scaling.""" + + def __init__(self, c: int, dilation: int = 1): + super().__init__() + self.conv = nn.Conv2d(c, c, 3, 1, dilation, dilation=dilation, groups=1) + self.beta = nn.Parameter(torch.ones((1, c, 1, 1)), requires_grad=True) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.relu(self.conv(x) * self.beta + x) + + +class IFBlock(nn.Module): + """Single-scale optical flow, mask, and feature block.""" + + def __init__(self, in_planes: int, c: int = 64): + super().__init__() + self.conv0 = nn.Sequential( + _conv(in_planes, c // 2, 3, 2, 1), + _conv(c // 2, c, 3, 2, 1), + ) + self.convblock = nn.Sequential( + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ResConv(c), + ) + self.lastconv = nn.Sequential( + nn.ConvTranspose2d(c, 4 * 13, 4, 2, 1), + nn.PixelShuffle(2), + ) + + def forward( + self, + x: torch.Tensor, + flow: torch.Tensor | None = None, + scale: float = 1.0, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + x = F.interpolate(x, scale_factor=1.0 / scale, mode="bilinear", align_corners=False) + if flow is not None: + flow = ( + F.interpolate( + flow, + scale_factor=1.0 / scale, + mode="bilinear", + align_corners=False, + ) + * 1.0 + / scale + ) + x = torch.cat((x, flow), 1) + feat = self.conv0(x) + feat = self.convblock(feat) + tmp = self.lastconv(feat) + tmp = F.interpolate(tmp, scale_factor=scale, mode="bilinear", align_corners=False) + flow = tmp[:, :4] * scale + mask = tmp[:, 4:5] + feat = tmp[:, 5:] + return flow, mask, feat + + +class Head(nn.Module): + """Feature encoder producing four-channel features at full resolution.""" + + def __init__(self): + super().__init__() + self.cnn0 = nn.Conv2d(3, 16, 3, 2, 1) + self.cnn1 = nn.Conv2d(16, 16, 3, 1, 1) + self.cnn2 = nn.Conv2d(16, 16, 3, 1, 1) + self.cnn3 = nn.ConvTranspose2d(16, 4, 4, 2, 1) + self.relu = nn.LeakyReLU(0.2, True) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x0 = self.cnn0(x) + x = self.relu(x0) + x1 = self.cnn1(x) + x = self.relu(x1) + x2 = self.cnn2(x) + x = self.relu(x2) + x3 = self.cnn3(x) + return x3 + + +class IFNet(nn.Module): + """Four-scale IFNet optical flow network.""" + + def __init__(self): + super().__init__() + self.block0 = IFBlock(7 + 8, c=192) + self.block1 = IFBlock(8 + 4 + 8 + 8, c=128) + self.block2 = IFBlock(8 + 4 + 8 + 8, c=64) + self.block3 = IFBlock(8 + 4 + 8 + 8, c=32) + self.encode = Head() + + def forward( + self, + x: torch.Tensor, + timestep: float = 0.5, + scale_list: list[float] | None = None, + ) -> tuple[list[torch.Tensor], torch.Tensor, list[tuple[torch.Tensor, torch.Tensor] | torch.Tensor]]: + if scale_list is None: + scale_list = [8, 4, 2, 1] + + channel = x.shape[1] // 2 + img0 = x[:, :channel] + img1 = x[:, channel:] + + if not torch.is_tensor(timestep): + timestep = (x[:, :1].clone() * 0 + 1) * timestep + else: + timestep = timestep.repeat(1, 1, img0.shape[2], img0.shape[3]) + + f0 = self.encode(img0[:, :3]) + f1 = self.encode(img1[:, :3]) + + flow_list: list[torch.Tensor] = [] + merged: list[tuple[torch.Tensor, torch.Tensor] | torch.Tensor] = [] + mask_list: list[torch.Tensor] = [] + warped_img0 = img0 + warped_img1 = img1 + flow = None + mask = None + + for i, block in enumerate([self.block0, self.block1, self.block2, self.block3]): + if flow is None: + flow, mask, feat = block( + torch.cat((img0[:, :3], img1[:, :3], f0, f1, timestep), 1), + None, + scale=scale_list[i], + ) + else: + wf0 = warp(f0, flow[:, :2]) + wf1 = warp(f1, flow[:, 2:4]) + fd, m0, feat = block( + torch.cat( + ( + warped_img0[:, :3], + warped_img1[:, :3], + wf0, + wf1, + timestep, + mask, + feat, + ), + 1, + ), + flow, + scale=scale_list[i], + ) + mask = m0 + flow = flow + fd + + mask_list.append(mask) + flow_list.append(flow) + warped_img0 = warp(img0, flow[:, :2]) + warped_img1 = warp(img1, flow[:, 2:4]) + merged.append((warped_img0, warped_img1)) + + mask = torch.sigmoid(mask) + merged[3] = warped_img0 * mask + warped_img1 * (1 - mask) + return flow_list, mask_list[3], merged + + +class Model: + """Wraps IFNet and exposes RIFE-compatible load/inference helpers.""" + + def __init__(self): + self.flownet = IFNet() + + def eval(self) -> Model: + self.flownet.eval() + return self + + def device(self) -> torch.device: + return next(self.flownet.parameters()).device + + def load_model(self, path: str) -> None: + flownet_path = os.path.join(path, "flownet.pkl") + if not os.path.isfile(flownet_path): + raise FileNotFoundError( + f"RIFE weight file not found: {flownet_path}. Expected layout: /flownet.pkl" + ) + + state = torch.load(flownet_path, map_location="cpu", weights_only=False) + state = {k.removeprefix("module."): v for k, v in state.items()} + self.flownet.load_state_dict(state, strict=False) + logger.info("Loaded RIFE weights from %s", flownet_path) + + def inference( + self, + img0: torch.Tensor, + img1: torch.Tensor, + scale: float = 1.0, + timestep: float = 0.5, + ) -> torch.Tensor: + _n, _c, h, w = img0.shape + ph = ((h - 1) // 32 + 1) * 32 + pw = ((w - 1) // 32 + 1) * 32 + pad = (0, pw - w, 0, ph - h) + img0 = F.pad(img0, pad) + img1 = F.pad(img1, pad) + + imgs = torch.cat((img0, img1), 1) + scale_list = [8 / scale, 4 / scale, 2 / scale, 1 / scale] + with torch.no_grad(): + _flow_list, _mask, merged = self.flownet( + imgs, + timestep=timestep, + scale_list=scale_list, + ) + return merged[3][:, :, :h, :w] + + +def _resolve_rife_model_path(model_path: str | None) -> str: + model_path = model_path or _DEFAULT_RIFE_HF_REPO + if os.path.isdir(model_path): + return model_path + from vllm_omni.model_executor.model_loader.weight_utils import ( + download_weights_from_hf_specific, + ) + + return download_weights_from_hf_specific( + model_path, + cache_dir=None, + allow_patterns=["flownet.pkl"], + require_all=True, + ) + + +def _select_torch_device() -> torch.device: + try: + from vllm_omni.platforms import current_omni_platform + + return current_omni_platform.get_torch_device() + except Exception as exc: + logger.warning("Failed to resolve current vLLM-Omni torch device: %s", exc) + + if torch.cuda.is_available(): + return torch.device("cuda") + return torch.device("cpu") + + +def _normalize_video_tensor_layout(video: torch.Tensor) -> tuple[torch.Tensor, Any]: + if video.ndim == 5: + if video.shape[1] in (3, 4): + return video, lambda out: out + if video.shape[2] in (3, 4): + return video.permute(0, 2, 1, 3, 4), lambda out: out.permute(0, 2, 1, 3, 4) + elif video.ndim == 4: + if video.shape[0] in (3, 4): + return video.unsqueeze(0), lambda out: out.squeeze(0) + if video.shape[1] in (3, 4): + return video.permute(1, 0, 2, 3).unsqueeze(0), lambda out: out.squeeze(0).permute(1, 0, 2, 3) + raise ValueError(f"Unsupported video tensor shape for interpolation: {tuple(video.shape)}") + + +def _normalize_video_tensor_range(video: torch.Tensor) -> tuple[torch.Tensor, Any]: + original_dtype = video.dtype + video = video.detach() + if video.is_floating_point(): + video = video.to(torch.float32) + if torch.amin(video) < 0.0 or torch.amax(video) > 1.0: + return video.clamp(-1.0, 1.0) * 0.5 + 0.5, lambda out: (out * 2.0 - 1.0).to(original_dtype) + return video.clamp(0.0, 1.0), lambda out: out.to(original_dtype) + return video.to(torch.float32) / 255.0, lambda out: (out * 255.0).round().clamp(0, 255).to(original_dtype) + + +class FrameInterpolator: + """Lazy-loaded RIFE 4.22.lite frame interpolator.""" + + def __init__(self, model_path: str | None = None): + self._model_path = model_path + self._resolved_path: str | None = None + + def _ensure_model_loaded(self, preferred_device: torch.device | None = None) -> Model: + resolved_path = _resolve_rife_model_path(self._model_path) + self._resolved_path = resolved_path + device = preferred_device or _select_torch_device() + cache_key = (resolved_path, str(device)) + + with _MODEL_CACHE_LOCK: + if cache_key in _MODEL_CACHE: + return _MODEL_CACHE[cache_key] + + model = Model() + model.load_model(resolved_path) + model.eval() + model.flownet = model.flownet.to(device) + _MODEL_CACHE[cache_key] = model + logger.info("RIFE model loaded on device: %s", device) + return model + + def _make_inference( + self, + model: Model, + img0: torch.Tensor, + img1: torch.Tensor, + n: int, + scale: float, + ) -> list[torch.Tensor]: + if n == 1: + return [model.inference(img0, img1, scale=scale)] + mid = model.inference(img0, img1, scale=scale) + return ( + self._make_inference(model, img0, mid, n // 2, scale) + + [mid] + + self._make_inference(model, mid, img1, n // 2, scale) + ) + + def interpolate_tensor( + self, + video: torch.Tensor, + exp: int = 1, + scale: float = 1.0, + ) -> tuple[torch.Tensor, int]: + if exp < 1: + raise ValueError(f"frame interpolation exp must be >= 1, got {exp}") + if scale <= 0: + raise ValueError(f"frame interpolation scale must be > 0, got {scale}") + + video, restore_layout = _normalize_video_tensor_layout(video) + if video.shape[2] < 2: + return restore_layout(video), 1 + + video, restore_range = _normalize_video_tensor_range(video) + # Prefer the decoded video's current device so CPU-offloaded requests do + # not move the tensor back to GPU just for interpolation. + model = self._ensure_model_loaded(preferred_device=video.device) + video = video.to(model.device()) + intermediates_per_pair = 2**exp // 2 + + result_frames: list[torch.Tensor] = [] + for idx in range(video.shape[2] - 1): + img0 = video[:, :, idx, :, :] + img1 = video[:, :, idx + 1, :, :] + result_frames.append(img0) + result_frames.extend(self._make_inference(model, img0, img1, intermediates_per_pair, scale)) + result_frames.append(video[:, :, -1, :, :]) + result = torch.stack(result_frames, dim=2) + return restore_layout(restore_range(result)), 2**exp + + +def interpolate_video_tensor( + video: torch.Tensor, + exp: int = 1, + scale: float = 1.0, + model_path: str | None = None, +) -> tuple[torch.Tensor, int]: + """Interpolate a video tensor and return the FPS multiplier.""" + interpolator = FrameInterpolator(model_path=model_path) + return interpolator.interpolate_tensor(video, exp=exp, scale=scale) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index d847a96db..11ba59e43 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -2075,6 +2075,10 @@ async def _parse_video_form( true_cfg_scale: float | None = Form(default=None), seed: int | None = Form(default=None), negative_prompt: str | None = Form(default=None), + enable_frame_interpolation: bool = Form(default=False), + frame_interpolation_exp: int = Form(default=1, ge=1), + frame_interpolation_scale: float = Form(default=1.0, gt=0.0), + frame_interpolation_model_path: str | None = Form(default=None), lora: str | None = Form(default=None), extra_params: str | None = Form(default=None), ) -> tuple[VideoGenerationRequest, "OmniOpenAIServingVideo", str, ReferenceImage | None]: @@ -2111,6 +2115,10 @@ async def _parse_video_form( "true_cfg_scale": true_cfg_scale, "seed": seed, "negative_prompt": negative_prompt, + "enable_frame_interpolation": enable_frame_interpolation, + "frame_interpolation_exp": frame_interpolation_exp, + "frame_interpolation_scale": frame_interpolation_scale, + "frame_interpolation_model_path": frame_interpolation_model_path, "lora": _parse_form_json(lora, expected_type=dict), "extra_params": _parse_form_json(extra_params, expected_type=dict), } diff --git a/vllm_omni/entrypoints/openai/protocol/videos.py b/vllm_omni/entrypoints/openai/protocol/videos.py index de5362dd9..7c2c3164d 100644 --- a/vllm_omni/entrypoints/openai/protocol/videos.py +++ b/vllm_omni/entrypoints/openai/protocol/videos.py @@ -150,6 +150,29 @@ class VideoGenerationRequest(BaseModel): ) seed: int | None = Field(default=None, description="Random seed for reproducibility") + # vllm-omni extensions for post-generation frame interpolation. + enable_frame_interpolation: bool = Field( + default=False, + description="Enable post-generation RIFE frame interpolation before MP4 encoding.", + ) + frame_interpolation_exp: int = Field( + default=1, + ge=1, + description="Interpolation exponent: 1=2x temporal resolution, 2=4x, etc.", + ) + frame_interpolation_scale: float = Field( + default=1.0, + gt=0.0, + description="RIFE inference scale. Use 0.5 for high-resolution inputs to save memory.", + ) + frame_interpolation_model_path: str | None = Field( + default=None, + description=( + "Local directory or Hugging Face repo ID containing RIFE flownet.pkl weights. " + "Defaults to elfgum/RIFE-4.22.lite." + ), + ) + # vllm-omni extension for per-request LoRA. lora: dict[str, Any] | None = Field( default=None, diff --git a/vllm_omni/entrypoints/openai/serving_video.py b/vllm_omni/entrypoints/openai/serving_video.py index 0001fa65f..741295c7c 100644 --- a/vllm_omni/entrypoints/openai/serving_video.py +++ b/vllm_omni/entrypoints/openai/serving_video.py @@ -113,6 +113,10 @@ class OmniOpenAIServingVideo: if vp.fps is not None: gen_params.fps = vp.fps gen_params.frame_rate = float(vp.fps) + gen_params.enable_frame_interpolation = request.enable_frame_interpolation + gen_params.frame_interpolation_exp = request.frame_interpolation_exp + gen_params.frame_interpolation_scale = request.frame_interpolation_scale + gen_params.frame_interpolation_model_path = request.frame_interpolation_model_path if request.num_inference_steps is not None: gen_params.num_inference_steps = request.num_inference_steps @@ -160,7 +164,7 @@ class OmniOpenAIServingVideo: videos = self._extract_video_outputs(result) audios = self._extract_audio_outputs(result, expected_count=len(videos)) audio_sample_rate = self._resolve_audio_sample_rate(result) - output_fps = vp.fps or self._resolve_fps(result) or 24 + output_fps = (vp.fps or self._resolve_fps(result) or 24) * self._resolve_video_fps_multiplier(result) return VideoGenerationArtifacts( videos=videos, audios=audios, @@ -243,6 +247,22 @@ class OmniOpenAIServingVideo: logger.info("Video response encoding (MP4 bytes): %.2f ms", _t_encode_ms) return video_bytes, artifacts.stage_durations, artifacts.peak_memory_mb + @staticmethod + def _resolve_video_fps_multiplier(result: Any) -> int: + custom_output = getattr(result, "custom_output", None) + if isinstance(custom_output, dict): + multiplier = custom_output.get("video_fps_multiplier") + if multiplier is not None: + return int(multiplier) + request_output = getattr(result, "request_output", None) + if request_output is not None: + custom_output = getattr(request_output, "custom_output", None) + if isinstance(custom_output, dict): + multiplier = custom_output.get("video_fps_multiplier") + if multiplier is not None: + return int(multiplier) + return 1 + @staticmethod def _apply_lora(lora_body: Any, gen_params: OmniDiffusionSamplingParams) -> None: try: diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 9cb6c4433..85faf6b94 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -227,6 +227,10 @@ class OmniDiffusionSamplingParams: frame_rate: float | None = None # Floating-point rate used by the diffusion model when it differs from `fps`. height_not_provided: bool = False width_not_provided: bool = False + enable_frame_interpolation: bool = False + frame_interpolation_exp: int = 1 + frame_interpolation_scale: float = 1.0 + frame_interpolation_model_path: str | None = None # Timesteps timesteps: torch.Tensor | None = None From 1ad726f49524be5a4fb96f777ed90722f1276692 Mon Sep 17 00:00:00 2001 From: TaffyOfficial <2587297563@qq.com> Date: Wed, 15 Apr 2026 12:53:47 +0800 Subject: [PATCH 08/97] =?UTF-8?q?[Fix]=20HunyuanImage-3.0:=20unify=20namin?= =?UTF-8?q?g=20hunyuan=5Fimage=5F3=20=E2=86=92=20hunyuan=5Fimage3=20(#2712?= =?UTF-8?q?)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- docs/design/feature/expert_parallel.md | 4 ++-- .../test_hunyuan_fused_moe.py | 22 +++++++++---------- .../test_hunyuanimage3_text2img.py | 2 +- .../__init__.py | 6 ++--- .../autoencoder.py | 0 .../hunyuan_fused_moe.py | 0 .../hunyuan_image3_tokenizer.py} | 2 +- .../hunyuan_image3_transformer.py} | 2 +- .../pipeline_hunyuan_image3.py} | 4 ++-- .../system_prompt.py | 0 vllm_omni/diffusion/registry.py | 4 ++-- ...age_3_moe.yaml => hunyuan_image3_moe.yaml} | 0 ...3_moe_dit.yaml => hunyuan_image3_t2i.yaml} | 0 ...2gpu.yaml => hunyuan_image3_t2i_2gpu.yaml} | 0 vllm_omni/platforms/interface.py | 2 +- vllm_omni/platforms/musa/platform.py | 2 +- ...3_moe_dit.yaml => hunyuan_image3_t2i.yaml} | 0 ...age_3_moe.yaml => hunyuan_image3_t2i.yaml} | 0 18 files changed, 25 insertions(+), 25 deletions(-) rename tests/diffusion/models/{hunyuan_image_3 => hunyuan_image3}/test_hunyuan_fused_moe.py (85%) rename vllm_omni/diffusion/models/{hunyuan_image_3 => hunyuan_image3}/__init__.py (58%) rename vllm_omni/diffusion/models/{hunyuan_image_3 => hunyuan_image3}/autoencoder.py (100%) rename vllm_omni/diffusion/models/{hunyuan_image_3 => hunyuan_image3}/hunyuan_fused_moe.py (100%) rename vllm_omni/diffusion/models/{hunyuan_image_3/hunyuan_image_3_tokenizer.py => hunyuan_image3/hunyuan_image3_tokenizer.py} (99%) rename vllm_omni/diffusion/models/{hunyuan_image_3/hunyuan_image_3_transformer.py => hunyuan_image3/hunyuan_image3_transformer.py} (99%) rename vllm_omni/diffusion/models/{hunyuan_image_3/pipeline_hunyuan_image_3.py => hunyuan_image3/pipeline_hunyuan_image3.py} (99%) rename vllm_omni/diffusion/models/{hunyuan_image_3 => hunyuan_image3}/system_prompt.py (100%) rename vllm_omni/model_executor/stage_configs/{hunyuan_image_3_moe.yaml => hunyuan_image3_moe.yaml} (100%) rename vllm_omni/model_executor/stage_configs/{hunyuan_image3_moe_dit.yaml => hunyuan_image3_t2i.yaml} (100%) rename vllm_omni/model_executor/stage_configs/{hunyuan_image_3_moe_2gpu.yaml => hunyuan_image3_t2i_2gpu.yaml} (100%) rename vllm_omni/platforms/npu/stage_configs/{hunyuan_image3_moe_dit.yaml => hunyuan_image3_t2i.yaml} (100%) rename vllm_omni/platforms/xpu/stage_configs/{hunyuan_image_3_moe.yaml => hunyuan_image3_t2i.yaml} (100%) diff --git a/docs/design/feature/expert_parallel.md b/docs/design/feature/expert_parallel.md index 9a7c4cdba..e05eec336 100644 --- a/docs/design/feature/expert_parallel.md +++ b/docs/design/feature/expert_parallel.md @@ -207,9 +207,9 @@ Complete examples in the codebase: | Model | Path | Pattern | Notes | |-------|------|---------|-------| -| **HunyuanImage3.0** | `vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py` | Standard EP | Full implementation with validation | +| **HunyuanImage3.0** | `vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py` | Standard EP | Full implementation with validation | | **EP Tests** | `vllm-omni/tests/e2e/offline_inference/test_expert_parallel.py` | E2E testing | EP correctness and performance | -| **Constraint Tests** | `vllm-omni/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py` | Unit testing | Validation logic | +| **Constraint Tests** | `vllm-omni/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py` | Unit testing | Validation logic | --- ## Summary diff --git a/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py similarity index 85% rename from tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py rename to tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py index 2cda9116c..626f78eed 100644 --- a/tests/diffusion/models/hunyuan_image_3/test_hunyuan_fused_moe.py +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_fused_moe.py @@ -12,7 +12,7 @@ class TestSetForwardContextNumTokens: def test_sets_num_tokens_when_context_available(self, mocker): """num_tokens should be set on ForwardContext when available.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe mock_ctx = mocker.MagicMock() del mock_ctx.in_profile_run # simulate missing attr @@ -26,7 +26,7 @@ class TestSetForwardContextNumTokens: def test_sets_in_profile_run_only_if_missing(self, mocker): """in_profile_run should not be overwritten if already set.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe mock_ctx = mocker.MagicMock() mock_ctx.in_profile_run = True # already set @@ -40,7 +40,7 @@ class TestSetForwardContextNumTokens: def test_noop_when_context_unavailable(self, mocker): """Should do nothing when ForwardContext is not available.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe mocker.patch.object(hunyuan_moe._vllm_fc, "is_forward_context_available", return_value=False) mock_get = mocker.patch.object(hunyuan_moe._vllm_fc, "get_forward_context") @@ -55,11 +55,11 @@ class TestHunyuanFusedMoEPlatformDispatch: def test_default_platform_uses_default_impl_qualname(self, mocker): """HunyuanFusedMoE should resolve the impl class from the platform hook.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe mock_platform = mocker.MagicMock() mock_platform.get_diffusion_model_impl_qualname.return_value = ( - "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault" + "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault" ) mocker.patch.object( @@ -71,7 +71,7 @@ class TestHunyuanFusedMoEPlatformDispatch: mock_impl = mocker.MagicMock() mock_resolve.return_value = mock_impl - from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import ( + from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import ( HunyuanFusedMoE, ) @@ -80,7 +80,7 @@ class TestHunyuanFusedMoEPlatformDispatch: mock_platform.prepare_diffusion_op_runtime.assert_called_once_with("hunyuan_fused_moe") mock_platform.get_diffusion_model_impl_qualname.assert_called_once_with("hunyuan_fused_moe") mock_resolve.assert_called_once_with( - "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault" + "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault" ) mock_impl.assert_called_once_with(prefix="") @@ -90,7 +90,7 @@ class TestHunyuanFusedMoEFactory: def test_new_delegates_to_impl_class(self, mocker): """HunyuanFusedMoE(prefix=..., **kwargs) should instantiate and return impl instance.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe class MockImpl: def __init__(self, *, prefix: str = "", **kwargs): @@ -104,7 +104,7 @@ class TestHunyuanFusedMoEFactory: mock_impl_class = mocker.MagicMock(return_value=MockImpl(prefix="test", a=1)) mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class) - from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import ( + from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import ( HunyuanFusedMoE, ) @@ -119,7 +119,7 @@ class TestHunyuanFusedMoEFactory: def test_make_expert_params_mapping_delegates_to_impl(self, mocker): """make_expert_params_mapping should delegate to impl class method.""" - import vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe as hunyuan_moe + import vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe as hunyuan_moe expected_mapping = [("a", "b", 0, "c")] mock_platform = mocker.MagicMock() @@ -130,7 +130,7 @@ class TestHunyuanFusedMoEFactory: mock_impl_class.make_expert_params_mapping = mocker.MagicMock(return_value=expected_mapping) mocker.patch.object(hunyuan_moe, "resolve_obj_by_qualname", return_value=mock_impl_class) - from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import ( + from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import ( HunyuanFusedMoE, ) diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py b/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py index 79bb64dca..6898763e4 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py @@ -17,7 +17,7 @@ PROMPT = "A brown and white dog is running on the grass" MODEL_NAME = "tencent/HunyuanImage-3.0" LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32" REPO_ROOT = Path(__file__).resolve().parents[3] -STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image_3_moe.yaml" +STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_moe.yaml" pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py b/vllm_omni/diffusion/models/hunyuan_image3/__init__.py similarity index 58% rename from vllm_omni/diffusion/models/hunyuan_image_3/__init__.py rename to vllm_omni/diffusion/models/hunyuan_image3/__init__.py index cbc6a8ad1..6612bd855 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/__init__.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/__init__.py @@ -2,12 +2,12 @@ # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Hunyuan Image 3 diffusion model components.""" -from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import HunyuanFusedMoE -from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_image_3_transformer import ( +from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE +from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_image3_transformer import ( HunyuanImage3Model, HunyuanImage3Text2ImagePipeline, ) -from vllm_omni.diffusion.models.hunyuan_image_3.pipeline_hunyuan_image_3 import ( +from vllm_omni.diffusion.models.hunyuan_image3.pipeline_hunyuan_image3 import ( HunyuanImage3Pipeline, ) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/autoencoder.py b/vllm_omni/diffusion/models/hunyuan_image3/autoencoder.py similarity index 100% rename from vllm_omni/diffusion/models/hunyuan_image_3/autoencoder.py rename to vllm_omni/diffusion/models/hunyuan_image3/autoencoder.py diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_fused_moe.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py similarity index 100% rename from vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_fused_moe.py rename to vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py similarity index 99% rename from vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py rename to vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py index ce563f711..4a29e9df9 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_tokenizer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_tokenizer.py @@ -13,7 +13,7 @@ from diffusers.utils.outputs import BaseOutput from transformers import AutoTokenizer from vllm.logger import init_logger -from .hunyuan_image_3_transformer import ImageInfo, JointImageInfo, default +from .hunyuan_image3_transformer import ImageInfo, JointImageInfo, default logger = init_logger(__name__) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py similarity index 99% rename from vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py rename to vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index bc81ca9c3..327260ee0 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/hunyuan_image_3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -74,7 +74,7 @@ from vllm_omni.diffusion.distributed.sp_plan import ( ) from vllm_omni.diffusion.distributed.utils import get_local_device from vllm_omni.diffusion.layers.rope import RotaryEmbedding -from vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe import HunyuanFusedMoE +from vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe import HunyuanFusedMoE logger = logging.getLogger(__name__) diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py similarity index 99% rename from vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py rename to vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 7e9e2d278..2f140b48f 100644 --- a/vllm_omni/diffusion/models/hunyuan_image_3/pipeline_hunyuan_image_3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -25,8 +25,8 @@ from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPi from vllm_omni.diffusion.request import OmniDiffusionRequest from .autoencoder import AutoencoderKLConv3D -from .hunyuan_image_3_tokenizer import TokenizerWrapper -from .hunyuan_image_3_transformer import ( +from .hunyuan_image3_tokenizer import TokenizerWrapper +from .hunyuan_image3_transformer import ( CausalMMOutputWithPast, HunyuanImage3ImageProcessor, HunyuanImage3Model, diff --git a/vllm_omni/diffusion/models/hunyuan_image_3/system_prompt.py b/vllm_omni/diffusion/models/hunyuan_image3/system_prompt.py similarity index 100% rename from vllm_omni/diffusion/models/hunyuan_image_3/system_prompt.py rename to vllm_omni/diffusion/models/hunyuan_image3/system_prompt.py diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 97bc7fa29..517b061ec 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -119,8 +119,8 @@ _DIFFUSION_MODELS = { "FluxKontextPipeline", ), "HunyuanImage3ForCausalMM": ( - "hunyuan_image_3", - "pipeline_hunyuan_image_3", + "hunyuan_image3", + "pipeline_hunyuan_image3", "HunyuanImage3Pipeline", ), "Flux2KleinPipeline": ( diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe.yaml similarity index 100% rename from vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe.yaml rename to vllm_omni/model_executor/stage_configs/hunyuan_image3_moe.yaml diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_t2i.yaml similarity index 100% rename from vllm_omni/model_executor/stage_configs/hunyuan_image3_moe_dit.yaml rename to vllm_omni/model_executor/stage_configs/hunyuan_image3_t2i.yaml diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe_2gpu.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_t2i_2gpu.yaml similarity index 100% rename from vllm_omni/model_executor/stage_configs/hunyuan_image_3_moe_2gpu.yaml rename to vllm_omni/model_executor/stage_configs/hunyuan_image3_t2i_2gpu.yaml diff --git a/vllm_omni/platforms/interface.py b/vllm_omni/platforms/interface.py index 8f1e66747..b69731a67 100644 --- a/vllm_omni/platforms/interface.py +++ b/vllm_omni/platforms/interface.py @@ -64,7 +64,7 @@ class OmniPlatform(Platform): @classmethod def get_diffusion_model_impl_qualname(cls, op_name: str) -> str: if op_name == "hunyuan_fused_moe": - return "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault" + return "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault" raise NotImplementedError(f"Unsupported diffusion model op: {op_name}") @classmethod diff --git a/vllm_omni/platforms/musa/platform.py b/vllm_omni/platforms/musa/platform.py index fe1ccc6d0..64a70a9be 100644 --- a/vllm_omni/platforms/musa/platform.py +++ b/vllm_omni/platforms/musa/platform.py @@ -39,7 +39,7 @@ class MUSAOmniPlatform(OmniPlatform, MUSAPlatformBase): def get_diffusion_model_impl_qualname(cls, op_name: str) -> str: # MUSA uses default implementations for diffusion ops if op_name == "hunyuan_fused_moe": - return "vllm_omni.diffusion.models.hunyuan_image_3.hunyuan_fused_moe.HunyuanFusedMoEDefault" + return "vllm_omni.diffusion.models.hunyuan_image3.hunyuan_fused_moe.HunyuanFusedMoEDefault" return super().get_diffusion_model_impl_qualname(op_name) @classmethod diff --git a/vllm_omni/platforms/npu/stage_configs/hunyuan_image3_moe_dit.yaml b/vllm_omni/platforms/npu/stage_configs/hunyuan_image3_t2i.yaml similarity index 100% rename from vllm_omni/platforms/npu/stage_configs/hunyuan_image3_moe_dit.yaml rename to vllm_omni/platforms/npu/stage_configs/hunyuan_image3_t2i.yaml diff --git a/vllm_omni/platforms/xpu/stage_configs/hunyuan_image_3_moe.yaml b/vllm_omni/platforms/xpu/stage_configs/hunyuan_image3_t2i.yaml similarity index 100% rename from vllm_omni/platforms/xpu/stage_configs/hunyuan_image_3_moe.yaml rename to vllm_omni/platforms/xpu/stage_configs/hunyuan_image3_t2i.yaml From 2dff2d7c747864378764195e0e4a6b137c3cf5df Mon Sep 17 00:00:00 2001 From: fan2956 Date: Wed, 15 Apr 2026 14:02:41 +0800 Subject: [PATCH 09/97] [PERF] Wan2.2 support adalayernorm fused op (#2585) Signed-off-by: fan2956 Co-authored-by: Canlin Guo --- .../diffusion/cache/teacache/extractors.py | 3 +- vllm_omni/diffusion/layers/adalayernorm.py | 80 +++++-------------- .../qwen_image/qwen_image_transformer.py | 19 +++-- .../models/wan2_2/wan2_2_transformer.py | 25 +++--- .../models/wan2_2/wan2_2_vace_transformer.py | 2 +- 5 files changed, 45 insertions(+), 84 deletions(-) diff --git a/vllm_omni/diffusion/cache/teacache/extractors.py b/vllm_omni/diffusion/cache/teacache/extractors.py index 3d247e318..84c237b60 100644 --- a/vllm_omni/diffusion/cache/teacache/extractors.py +++ b/vllm_omni/diffusion/cache/teacache/extractors.py @@ -222,7 +222,8 @@ def extract_qwen_context( block = module.transformer_blocks[0] img_mod_params = block.img_mod(temb) img_mod1, _ = img_mod_params.chunk(2, dim=-1) - img_modulated, _ = block.img_norm1(hidden_states, img_mod1) + img_scale1, img_shift1, _ = block._modulate(img_mod1) + img_modulated = block.img_norm1(hidden_states, img_scale1, img_shift1) # ============================================================================ # DEFINE TRANSFORMER EXECUTION (Qwen-specific) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 35f63e2fc..4d70ed52f 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -29,105 +29,61 @@ class AdaLayerNorm(CustomOp): self.hidden_size = hidden_size self.layernorm = nn.LayerNorm(self.hidden_size, elementwise_affine=self.elementwise_affine, eps=self.eps) - def preprocess( - self, - mod_params: torch.Tensor, - index: torch.Tensor = None, - ) -> torch.Tensor: - # shift: b d, scale: b d, gate: b d - shift, scale, gate = mod_params.chunk(3, dim=-1) - - if index is not None: - # Assuming mod_params batch dim is 2*actual_batch (chunked into 2 parts) - # So shift, scale, gate have shape [2*actual_batch, d] - actual_batch = shift.size(0) // 2 - shift_0, shift_1 = shift[:actual_batch], shift[actual_batch:] # each: [actual_batch, d] - scale_0, scale_1 = scale[:actual_batch], scale[actual_batch:] - gate_0, gate_1 = gate[:actual_batch], gate[actual_batch:] - - # index: [b, l] where b is actual batch size - # Expand to [b, l, 1] to match feature dimension - index_expanded = index.unsqueeze(-1) # [b, l, 1] - - # Expand chunks to [b, 1, d] then broadcast to [b, l, d] - shift_0_exp = shift_0.unsqueeze(1) # [b, 1, d] - shift_1_exp = shift_1.unsqueeze(1) # [b, 1, d] - scale_0_exp = scale_0.unsqueeze(1) - scale_1_exp = scale_1.unsqueeze(1) - gate_0_exp = gate_0.unsqueeze(1) - gate_1_exp = gate_1.unsqueeze(1) - - # Use torch.where to select based on index - shift_result = torch.where(index_expanded == 0, shift_0_exp, shift_1_exp) - scale_result = torch.where(index_expanded == 0, scale_0_exp, scale_1_exp) - gate_result = torch.where(index_expanded == 0, gate_0_exp, gate_1_exp) - else: - shift_result = shift.unsqueeze(1) - scale_result = scale.unsqueeze(1) - gate_result = gate.unsqueeze(1) - - return shift_result, scale_result, gate_result - def forward_cuda( self, x: torch.Tensor, - mod_params: torch.Tensor, - index: torch.Tensor = None, + scale: torch.Tensor, + shift: torch.Tensor, ) -> torch.Tensor: - return self.forward_native(x, mod_params, index) + return self.forward_native(x, scale, shift) def forward_hip( self, x: torch.Tensor, - mod_params: torch.Tensor, - index: torch.Tensor = None, + scale: torch.Tensor, + shift: torch.Tensor, ) -> torch.Tensor: - return self.forward_native(x, mod_params, index) + return self.forward_native(x, scale, shift) def forward_npu( self, x: torch.Tensor, - mod_params: torch.Tensor, - index: torch.Tensor = None, + scale: torch.Tensor, + shift: torch.Tensor, ) -> torch.Tensor: - shift_result, scale_result, gate_result = self.preprocess(mod_params, index) - if _HAS_MINDIESD: try: from mindiesd import layernorm_scale_shift - output = layernorm_scale_shift(self.layernorm, x, scale_result, shift_result, fused=True) + output = layernorm_scale_shift(self.layernorm, x, scale, shift, fused=True) - return output, gate_result + return output except ImportError as e: logger.warning_once(f"mindiesd import failed, falling back to torch_npu: {e}") import torch_npu output = ( - torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) * (1 + scale_result) - + shift_result + torch_npu.npu_layer_norm_eval(x, normalized_shape=[self.hidden_size], eps=self.eps) * (1 + scale) + shift ) - return output, gate_result + return output def forward_xpu( self, x: torch.Tensor, - mod_params: torch.Tensor, - index: torch.Tensor = None, + scale: torch.Tensor, + shift: torch.Tensor, ) -> torch.Tensor: - return self.forward_native(x, mod_params, index) + return self.forward_native(x, scale, shift) def forward_native( self, x: torch.Tensor, - mod_params: torch.Tensor, - index: torch.Tensor = None, + scale: torch.Tensor, + shift: torch.Tensor, ) -> torch.Tensor: - shift_result, scale_result, gate_result = self.preprocess(mod_params, index) - - return self.layernorm(x) * (1 + scale_result) + shift_result, gate_result + return self.layernorm(x) * (1 + scale) + shift class AdaLayerNormZero(nn.Module): diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index b34f19e95..9f16d8808 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -744,9 +744,9 @@ class QwenImageTransformerBlock(nn.Module): self.zero_cond_t = zero_cond_t - def _modulate(self, x, mod_params, index=None): + def _modulate(self, mod_params, index=None): """Apply modulation to input tensor""" - # x: b l d, shift: b d, scale: b d, gate: b d + # shift: b d, scale: b d, gate: b d shift, scale, gate = mod_params.chunk(3, dim=-1) if index is not None: @@ -778,7 +778,7 @@ class QwenImageTransformerBlock(nn.Module): scale_result = scale.unsqueeze(1) gate_result = gate.unsqueeze(1) - return x * (1 + scale_result) + shift_result, gate_result + return scale_result, shift_result, gate_result def forward( self, @@ -804,10 +804,12 @@ class QwenImageTransformerBlock(nn.Module): txt_mod1, txt_mod2 = txt_mod_params.chunk(2, dim=-1) # Each [B, 3*dim] # Process image stream - norm1 + modulation - img_modulated, img_gate1 = self.img_norm1(hidden_states, img_mod1, modulate_index) + img_scale1, img_shift1, img_gate1 = self._modulate(img_mod1, modulate_index) + img_modulated = self.img_norm1(hidden_states, img_scale1, img_shift1) # Process text stream - norm1 + modulation - txt_modulated, txt_gate1 = self.txt_norm1(encoder_hidden_states, txt_mod1) + txt_scale1, txt_shift1, txt_gate1 = self._modulate(txt_mod1) + txt_modulated = self.txt_norm1(encoder_hidden_states, txt_scale1, txt_shift1) # Use QwenAttnProcessor2_0 for joint attention computation # This directly implements the DoubleStreamLayerMegatron logic: @@ -832,13 +834,16 @@ class QwenImageTransformerBlock(nn.Module): encoder_hidden_states = encoder_hidden_states + txt_gate1 * txt_attn_output # Process image stream - norm2 + MLP - img_modulated2, img_gate2 = self.img_norm2(hidden_states, img_mod2, modulate_index) + img_scale2, img_shift2, img_gate2 = self._modulate(img_mod2, modulate_index) + img_modulated2 = self.img_norm2(hidden_states, img_scale2, img_shift2) img_mlp_output = self.img_mlp(img_modulated2) hidden_states = hidden_states + img_gate2 * img_mlp_output # Process text stream - norm2 + MLP - txt_modulated2, txt_gate2 = self.txt_norm2(encoder_hidden_states, txt_mod2) + txt_scale2, txt_shift2, txt_gate2 = self._modulate(txt_mod2) + txt_modulated2 = self.txt_norm2(encoder_hidden_states, txt_scale2, txt_shift2) + txt_mlp_output = self.txt_mlp(txt_modulated2) encoder_hidden_states = encoder_hidden_states + txt_gate2 * txt_mlp_output diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index 3b43f3eaf..b870193a1 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -29,6 +29,7 @@ from vllm_omni.diffusion.distributed.sp_plan import ( SequenceParallelOutput, ) from vllm_omni.diffusion.forward_context import get_forward_context +from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -620,7 +621,7 @@ class WanTransformerBlock(nn.Module): head_dim = dim // num_heads # 1. Self-attention - self.norm1 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.norm1 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) self.attn1 = WanSelfAttention( dim=dim, num_heads=num_heads, @@ -640,7 +641,7 @@ class WanTransformerBlock(nn.Module): # 3. Feed-forward self.ffn = WanFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim) - self.norm3 = FP32LayerNorm(dim, eps, elementwise_affine=False) + self.norm3 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) # Scale-shift table for modulation self.scale_shift_table = nn.Parameter(torch.randn(1, 6, dim) / dim**0.5) @@ -656,7 +657,7 @@ class WanTransformerBlock(nn.Module): if temb.ndim == 4: # temb: batch_size, seq_len, 6, inner_dim (wan2.2 ti2v) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table.unsqueeze(0) + temb.float() + self.scale_shift_table.unsqueeze(0) + temb ).chunk(6, dim=2) shift_msa = shift_msa.squeeze(2) scale_msa = scale_msa.squeeze(2) @@ -667,25 +668,23 @@ class WanTransformerBlock(nn.Module): else: # temb: batch_size, 6, inner_dim (wan2.1/wan2.2 14B) shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = ( - self.scale_shift_table + temb.float() + self.scale_shift_table + temb ).chunk(6, dim=1) # 1. Self-attention - norm_hidden_states = (self.norm1(hidden_states.float()) * (1 + scale_msa) + shift_msa).type_as(hidden_states) + norm_hidden_states = self.norm1(hidden_states, scale_msa, shift_msa).type_as(hidden_states) attn_output = self.attn1(norm_hidden_states, rotary_emb, hidden_states_mask) - hidden_states = (hidden_states.float() + attn_output * gate_msa).type_as(hidden_states) + hidden_states = (hidden_states + attn_output * gate_msa).type_as(hidden_states) # 2. Cross-attention - norm_hidden_states = self.norm2(hidden_states.float()).type_as(hidden_states) + norm_hidden_states = self.norm2(hidden_states).type_as(hidden_states) attn_output = self.attn2(norm_hidden_states, encoder_hidden_states) hidden_states = hidden_states + attn_output # 3. Feed-forward - norm_hidden_states = (self.norm3(hidden_states.float()) * (1 + c_scale_msa) + c_shift_msa).type_as( - hidden_states - ) + norm_hidden_states = self.norm3(hidden_states, c_scale_msa, c_shift_msa).type_as(hidden_states) ff_output = self.ffn(norm_hidden_states) - hidden_states = (hidden_states.float() + ff_output.float() * c_gate_msa).type_as(hidden_states) + hidden_states = (hidden_states + ff_output * c_gate_msa).type_as(hidden_states) return hidden_states @@ -854,7 +853,7 @@ class WanTransformer3DModel(nn.Module): ) # 4. Output norm & projection - self.norm_out = FP32LayerNorm(inner_dim, eps, elementwise_affine=False) + self.norm_out = AdaLayerNorm(inner_dim, elementwise_affine=False, eps=eps) self.proj_out = nn.Linear(inner_dim, out_channels * math.prod(patch_size)) # SP helper modules @@ -942,7 +941,7 @@ class WanTransformer3DModel(nn.Module): shift = shift.unsqueeze(1) scale = scale.unsqueeze(1) - hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) + hidden_states = self.norm_out(hidden_states, scale, shift).type_as(hidden_states) hidden_states = self.proj_out(hidden_states) hidden_states = hidden_states.reshape( diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py index 4f4217dab..c48938e1b 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_vace_transformer.py @@ -239,7 +239,7 @@ class WanVACETransformer3DModel(WanTransformer3DModel): shift = shift.unsqueeze(1) scale = scale.unsqueeze(1) - hidden_states = (self.norm_out(hidden_states.float()) * (1 + scale) + shift).type_as(hidden_states) + hidden_states = self.norm_out(hidden_states, scale, shift).type_as(hidden_states) hidden_states = self.proj_out(hidden_states) hidden_states = hidden_states.reshape( From 133e2f97068f4ae57fc91d7afd1e405386a0e12e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Zeyu=20Huang=20=7C=20=E9=BB=83=E6=BE=A4=E5=AE=87?= <11222265+fhfuih@users.noreply.github.com> Date: Wed, 15 Apr 2026 14:08:00 +0800 Subject: [PATCH 10/97] [hotfix] API connection error in CI (#2810) --- tests/conftest.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/conftest.py b/tests/conftest.py index 4ad4706fc..098fd8d97 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -2876,7 +2876,7 @@ class OpenAIClientHandler: return f"{self.base_url.rstrip('/')}/{path.lstrip('/')}" -@pytest.fixture(scope="module") +@pytest.fixture def openai_client(request: pytest.FixtureRequest, run_level: str): """Create OpenAIClientHandler fixture to facilitate communication with OmniServer with encapsulated request sending, concurrent requests, response handling, and validation.""" From 38d5f2d530c84cdb5462116103944b2b84e44182 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Wed, 15 Apr 2026 14:09:22 +0800 Subject: [PATCH 11/97] [Perf] VoxCPM2: Speedup by manual CUDA Graph capture for scaffold/residual forward (#2803) Signed-off-by: Sy03 <1370724210@qq.com> --- .../models/voxcpm2/minicpm4_paged.py | 20 ++ .../models/voxcpm2/voxcpm2_talker.py | 188 ++++++++++++++++-- 2 files changed, 189 insertions(+), 19 deletions(-) diff --git a/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py b/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py index 40bacfff6..b87ec5aaf 100644 --- a/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py +++ b/vllm_omni/model_executor/models/voxcpm2/minicpm4_paged.py @@ -307,6 +307,16 @@ class MiniCPM4PagedForVoxCPM2(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states + def precompute_fused_qkv(self) -> None: + """Materialize fused QKV weights before CUDA Graph capture.""" + for layer in self.layers: + attn = layer.self_attn + if attn._fused_qkv_weight is None: + attn._fused_qkv_weight = torch.cat( + [attn.q_proj.weight, attn.k_proj.weight, attn.v_proj.weight], + dim=0, + ).detach() + def compile_selective(self) -> list[str]: """Compile the full model forward as one graph. @@ -411,6 +421,16 @@ class MiniCPM4PagedResidualLM(nn.Module): hidden_states = self.norm(hidden_states) return hidden_states + def precompute_fused_qkv(self) -> None: + """Materialize fused QKV weights before CUDA Graph capture.""" + for layer in self.layers: + attn = layer.self_attn + if attn._fused_qkv_weight is None: + attn._fused_qkv_weight = torch.cat( + [attn.q_proj.weight, attn.k_proj.weight, attn.v_proj.weight], + dim=0, + ).detach() + def compile_selective(self) -> list[str]: """Compile the full residual model forward as one graph (same strategy as base_lm).""" if self._compiled_layers: diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 94f065890..02bcae821 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -10,6 +10,7 @@ Architecture: from __future__ import annotations +import copy import dataclasses import logging import os @@ -21,6 +22,7 @@ import librosa import torch import torch.nn as nn from vllm.config import VllmConfig +from vllm.forward_context import get_forward_context, override_forward_context from vllm.logger import init_logger from vllm.model_executor.models.utils import ( AutoWeightsLoader, @@ -101,6 +103,14 @@ class _RequestState: last_decoded_audio: torch.Tensor | None = None +@dataclasses.dataclass +class _CapturedGraph: + graph: torch.cuda.CUDAGraph + input_embeds: torch.Tensor + positions: torch.Tensor + output: torch.Tensor + + # =================================================================== # Profiling timer # =================================================================== @@ -336,6 +346,13 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._perf = _PerfTimer(enabled=_ENABLE_PROFILING) self._cfm_buffers: _CFMBufferManager | None = None + self._enable_cuda_graph = True + self._scaffold_graphs: dict[int, _CapturedGraph] = {} + self._residual_graphs: dict[int, _CapturedGraph] = {} + self._max_cached_graphs = self._max_batch_size + self._cuda_graph_pool: tuple | None = None + self._cuda_graph_warmup_steps = 0 + self._cuda_graph_warmup_threshold = 3 self._active_states: dict[str, _RequestState] = {} self._current_request_id: str | None = None @@ -483,19 +500,24 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): except Exception as e: logger.warning("torch.compile AudioVAE failed: %s", e) - if not getattr(self.model, "_selective_compiled", False): - try: - targets.extend(f"scaffold.{t}" for t in self.model.compile_selective()) - self.model._selective_compiled = True - except Exception as e: - logger.warning("scaffold compile failed: %s", e) + if not self._enable_cuda_graph: + if not getattr(self.model, "_selective_compiled", False): + try: + targets.extend(f"scaffold.{t}" for t in self.model.compile_selective()) + self.model._selective_compiled = True + except Exception as e: + logger.warning("scaffold compile failed: %s", e) - if not getattr(self.residual_model, "_selective_compiled", False): - try: - targets.extend(f"residual.{t}" for t in self.residual_model.compile_selective()) - self.residual_model._selective_compiled = True - except Exception as e: - logger.warning("residual compile failed: %s", e) + if not getattr(self.residual_model, "_selective_compiled", False): + try: + targets.extend(f"residual.{t}" for t in self.residual_model.compile_selective()) + self.residual_model._selective_compiled = True + except Exception as e: + logger.warning("residual compile failed: %s", e) + else: + self.model.precompute_fused_qkv() + self.residual_model.precompute_fused_qkv() + targets.append("scaffold+residual (CUDA Graph, skipping compile)") if not getattr(self, "_projections_compiled", False): try: @@ -518,6 +540,90 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): tts = self.tts return tts.stop_head(tts.stop_actn(tts.stop_proj(lm_h))) + def _get_cuda_graph_pool(self) -> tuple: + if self._cuda_graph_pool is None: + self._cuda_graph_pool = torch.cuda.graph_pool_handle() + return self._cuda_graph_pool + + @staticmethod + def _nullify_volatile_metadata(ctx: Any) -> Any: + """Set ``scheduler_metadata`` to None on all attention layers. + + This is the only tensor FA3 reallocates each step (variable shape). + All other metadata tensors are persistent model-runner buffers. + Setting it to None makes FA3 use default scheduling (~0.1ms cost). + """ + if not isinstance(ctx.attn_metadata, dict): + return ctx + + ctx = copy.copy(ctx) + new_meta: dict[str, Any] = {} + for layer_name, meta in ctx.attn_metadata.items(): + if getattr(meta, "scheduler_metadata", None) is not None: + meta = copy.copy(meta) + meta.scheduler_metadata = None + new_meta[layer_name] = meta + ctx.attn_metadata = new_meta + return ctx + + def _capture_graph( + self, + model: nn.Module, + batch_size: int, + label: str, + is_residual: bool = False, + ) -> _CapturedGraph: + """Capture a CUDA Graph for *model* at *batch_size*.""" + hidden_size = self.config.hidden_size + dtype = self._side_dtype + dev = torch.device(self._device) + pool = self._get_cuda_graph_pool() + + model.precompute_fused_qkv() + + g = _CapturedGraph( + graph=torch.cuda.CUDAGraph(), + input_embeds=torch.zeros(batch_size, hidden_size, device=dev, dtype=dtype), + positions=torch.zeros(batch_size, device=dev, dtype=torch.long), + output=torch.zeros(batch_size, hidden_size, device=dev, dtype=dtype), + ) + + if is_residual: + call_kwargs = dict(positions=g.positions, inputs_embeds=g.input_embeds) + else: + call_kwargs = dict(input_ids=None, positions=g.positions, inputs_embeds=g.input_embeds) + + ctx = get_forward_context() + patched_ctx = self._nullify_volatile_metadata(ctx) + + with override_forward_context(patched_ctx): + for _ in range(3): + _ = model(**call_kwargs) + + with torch.cuda.graph(g.graph, pool=pool): + g.output = model(**call_kwargs) + + logger.info("CUDA Graph captured for %s (batch_size=%d)", label, batch_size) + return g + + def _replay_graph( + self, + g: _CapturedGraph, + inputs_embeds: torch.Tensor, + positions: torch.Tensor, + batch_size: int, + ) -> torch.Tensor: + """Copy fresh inputs into static buffers, then replay. + + No metadata copy needed: persistent buffers (seq_lens, slot_mapping, + etc.) are updated in-place by the model runner. scheduler_metadata + was nullified at capture time so no kernel references it. + """ + g.input_embeds[:batch_size].copy_(inputs_embeds[:batch_size]) + g.positions[:batch_size].copy_(positions[:batch_size]) + g.graph.replay() + return g.output[:batch_size].clone() + # -------------------- vllm hooks -------------------- def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: @@ -534,12 +640,35 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._perf.start("forward_total") dev = input_ids.device - model_output = self.model(input_ids, positions, intermediate_tensors, inputs_embeds) - if isinstance(model_output, IntermediateTensors): - return model_output - scaffold_hidden = model_output - if isinstance(scaffold_hidden, tuple): - scaffold_hidden = scaffold_hidden[0] + num_reqs = len(self._pending_requests) + num_decode = sum(1 for _, is_p, _, n in self._pending_requests if not is_p and n == 1) + is_all_decode = num_decode == num_reqs and num_reqs > 0 + + tts_compiled = getattr(self.tts.feat_decoder.estimator, "_compiled", False) if self._tts is not None else False + graph_ready = tts_compiled and self._cuda_graph_warmup_steps >= self._cuda_graph_warmup_threshold + if num_decode > 0: + self._cuda_graph_warmup_steps += 1 + + can_use_graph = ( + self._enable_cuda_graph and graph_ready and intermediate_tensors is None and inputs_embeds is not None + ) + + if can_use_graph and is_all_decode and num_reqs <= self._max_cached_graphs: + self._perf.start("scaffold_fwd") + if num_reqs not in self._scaffold_graphs: + self._scaffold_graphs[num_reqs] = self._capture_graph(self.model, num_reqs, "scaffold") + scaffold_hidden = self._replay_graph(self._scaffold_graphs[num_reqs], inputs_embeds, positions, num_reqs) + self._perf.stop("scaffold_fwd") + + else: + self._perf.start("scaffold_fwd") + model_output = self.model(input_ids, positions, intermediate_tensors, inputs_embeds) + self._perf.stop("scaffold_fwd") + if isinstance(model_output, IntermediateTensors): + return model_output + scaffold_hidden = model_output + if isinstance(scaffold_hidden, tuple): + scaffold_hidden = scaffold_hidden[0] # Phase 1: per-request FSQ + residual input token_offset = 0 @@ -571,7 +700,28 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): if residual_inputs: batch_in = torch.cat(residual_inputs, dim=0) batch_pos = torch.cat(residual_positions, dim=0) - batch_out = self.residual_model(batch_pos, batch_in) + + residual_batch_size = batch_in.shape[0] + use_residual_graph = ( + self._enable_cuda_graph + and is_all_decode + and graph_ready + and residual_batch_size == num_reqs # 1 token per request + and residual_batch_size <= self._max_cached_graphs + ) + + self._perf.start("residual_fwd") + if use_residual_graph: + if residual_batch_size not in self._residual_graphs: + self._residual_graphs[residual_batch_size] = self._capture_graph( + self.residual_model, residual_batch_size, "residual", is_residual=True + ) + batch_out = self._replay_graph( + self._residual_graphs[residual_batch_size], batch_in, batch_pos, residual_batch_size + ) + else: + batch_out = self.residual_model(batch_pos, batch_in) + self._perf.stop("residual_fwd") # Phase 3: per-request LocDiT + update offset = 0 From 4bf4c6314741da606ff2b99efde5a83713cd8a22 Mon Sep 17 00:00:00 2001 From: IsleOfDawnlight Date: Wed, 15 Apr 2026 15:04:58 +0800 Subject: [PATCH 12/97] Add voxcpm model support. (#2467) Signed-off-by: Celeste-jq <591998922@qq.com> Signed-off-by: lyj-jjj Signed-off-by: IsleOfDawnlight Signed-off-by: Yueqian Lin Co-authored-by: Celeste-jq <591998922@qq.com> Co-authored-by: lyj-jjj Co-authored-by: Yueqian Lin --- .buildkite/test-ready.yml | 25 + benchmarks/voxcpm/README.md | 119 +++ .../voxcpm/vllm_omni/bench_tts_offline.py | 890 ++++++++++++++++++ .../voxcpm/vllm_omni/bench_tts_serve.py | 283 ++++++ .../voxcpm/vllm_omni/run_offline_matrix.py | 303 ++++++ examples/offline_inference/voxcpm/README.md | 123 +++ examples/offline_inference/voxcpm/end2end.py | 206 ++++ examples/online_serving/voxcpm/README.md | 166 ++++ .../voxcpm/openai_speech_client.py | 155 +++ examples/online_serving/voxcpm/run_server.sh | 38 + tests/e2e/offline_inference/test_voxcpm.py | 156 +++ tests/engine/test_arg_utils.py | 19 + .../openai_api/test_serving_speech_voxcpm.py | 143 +++ tests/entrypoints/test_utils.py | 33 + .../test_voxcpm_async_chunk.py | 87 ++ vllm_omni/engine/arg_utils.py | 3 + .../entrypoints/openai/serving_speech.py | 72 +- vllm_omni/model_executor/models/registry.py | 6 + .../model_executor/models/voxcpm/__init__.py | 7 + .../models/voxcpm/configuration_voxcpm.py | 3 + .../model_executor/models/voxcpm/voxcpm.py | 886 +++++++++++++++++ .../models/voxcpm/voxcpm_loader.py | 247 +++++ .../models/voxcpm/voxcpm_runtime_utils.py | 44 + .../models/voxcpm/voxcpm_stage_wrappers.py | 185 ++++ .../model_executor/stage_configs/voxcpm.yaml | 69 ++ .../stage_configs/voxcpm_async_chunk.yaml | 102 ++ .../stage_input_processors/voxcpm.py | 128 +++ .../platforms/npu/stage_configs/voxcpm.yaml | 67 ++ .../npu/stage_configs/voxcpm_async_chunk.yaml | 93 ++ .../transformers_utils/configs/__init__.py | 3 + .../transformers_utils/configs/voxcpm.py | 68 ++ 31 files changed, 4727 insertions(+), 2 deletions(-) create mode 100644 benchmarks/voxcpm/README.md create mode 100644 benchmarks/voxcpm/vllm_omni/bench_tts_offline.py create mode 100644 benchmarks/voxcpm/vllm_omni/bench_tts_serve.py create mode 100644 benchmarks/voxcpm/vllm_omni/run_offline_matrix.py create mode 100644 examples/offline_inference/voxcpm/README.md create mode 100644 examples/offline_inference/voxcpm/end2end.py create mode 100644 examples/online_serving/voxcpm/README.md create mode 100644 examples/online_serving/voxcpm/openai_speech_client.py create mode 100755 examples/online_serving/voxcpm/run_server.sh create mode 100644 tests/e2e/offline_inference/test_voxcpm.py create mode 100644 tests/entrypoints/openai_api/test_serving_speech_voxcpm.py create mode 100644 tests/model_executor/stage_input_processors/test_voxcpm_async_chunk.py create mode 100644 vllm_omni/model_executor/models/voxcpm/__init__.py create mode 100644 vllm_omni/model_executor/models/voxcpm/configuration_voxcpm.py create mode 100644 vllm_omni/model_executor/models/voxcpm/voxcpm.py create mode 100644 vllm_omni/model_executor/models/voxcpm/voxcpm_loader.py create mode 100644 vllm_omni/model_executor/models/voxcpm/voxcpm_runtime_utils.py create mode 100644 vllm_omni/model_executor/models/voxcpm/voxcpm_stage_wrappers.py create mode 100644 vllm_omni/model_executor/stage_configs/voxcpm.yaml create mode 100644 vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml create mode 100644 vllm_omni/model_executor/stage_input_processors/voxcpm.py create mode 100644 vllm_omni/platforms/npu/stage_configs/voxcpm.yaml create mode 100644 vllm_omni/platforms/npu/stage_configs/voxcpm_async_chunk.yaml create mode 100644 vllm_omni/transformers_utils/configs/voxcpm.py diff --git a/.buildkite/test-ready.yml b/.buildkite/test-ready.yml index 2f749f0ee..68f8e6152 100644 --- a/.buildkite/test-ready.yml +++ b/.buildkite/test-ready.yml @@ -295,6 +295,31 @@ steps: volumes: - "/fsx/hf_cache:/fsx/hf_cache" + - label: "VoxCPM E2E Test" + timeout_in_minutes: 20 + depends_on: upload-ready-pipeline + commands: + - | + timeout 20m bash -c ' + pip install voxcpm + export VLLM_LOGGING_LEVEL=DEBUG + export VLLM_WORKER_MULTIPROC_METHOD=spawn + pytest -s -v tests/e2e/offline_inference/test_voxcpm.py -m "core_model" --run-level "core_model" + ' + agents: + queue: "gpu_1_queue" + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + - label: "VoxCPM2 Native AR E2E Test" timeout_in_minutes: 20 depends_on: upload-ready-pipeline diff --git a/benchmarks/voxcpm/README.md b/benchmarks/voxcpm/README.md new file mode 100644 index 000000000..17f904101 --- /dev/null +++ b/benchmarks/voxcpm/README.md @@ -0,0 +1,119 @@ +# VoxCPM Benchmark + +This directory contains both: + +- online serving benchmark through the OpenAI-compatible `/v1/audio/speech` API +- offline benchmark for `Omni` / `AsyncOmni` +- full offline smoke-matrix orchestration + +Both benchmark paths report: + +- TTFP: time to first PCM packet +- E2E latency +- RTF: real-time factor (`e2e / audio_duration`) + +## Offline Benchmark + +Single offline benchmark run: + +```bash +python benchmarks/voxcpm/vllm_omni/bench_tts_offline.py \ + --model /path/to/voxcpm-model \ + --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm.yaml \ + --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni." \ + --warmup-runs 1 \ + --output-dir benchmarks/voxcpm/results/offline_single +``` + +Streaming offline benchmark: + +```bash +python benchmarks/voxcpm/vllm_omni/bench_tts_offline.py \ + --model /path/to/voxcpm-model \ + --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \ + --text "This is a split-stage VoxCPM streaming example running on vLLM Omni." \ + --warmup-runs 1 \ + --output-dir benchmarks/voxcpm/results/offline_streaming +``` + +Full fixed offline matrix, equivalent to the old `examples/offline_inference/voxcpm/test.py`: + +```bash +python benchmarks/voxcpm/vllm_omni/run_offline_matrix.py \ + --model /path/to/voxcpm-model \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript spoken in reference.wav." \ + --output-root benchmarks/voxcpm/results/offline_matrix +``` + +The full matrix covers both routes: + +- streaming: `voxcpm_async_chunk.yaml` +- sync: `voxcpm.yaml` + +And these six scenarios under each route: + +- warmup + single TTS +- warmup + single voice cloning +- warmup + batch TTS +- warmup + batch voice cloning +- cold single TTS +- cold single voice cloning + +`bench_tts_offline.py` itself no longer writes `summary.json` / `results.json`; it prints TTFP / RTF inline and saves generated WAV files only. The matrix runner keeps only per-case `run.log`. + +## Start the Server + +Async-chunk: + +```bash +vllm serve /path/to/voxcpm-model \ + --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \ + --trust-remote-code \ + --enforce-eager \ + --omni \ + --port 8091 +``` + +Non-streaming: + +```bash +vllm serve /path/to/voxcpm-model \ + --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm.yaml \ + --trust-remote-code \ + --enforce-eager \ + --omni \ + --port 8091 +``` + +## Run the Benchmark + +```bash +python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \ + --host 127.0.0.1 \ + --port 8091 \ + --num-prompts 20 \ + --max-concurrency 1 \ + --result-dir /tmp/voxcpm_bench +``` + +Voice cloning benchmark: + +```bash +python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \ + --host 127.0.0.1 \ + --port 8091 \ + --num-prompts 10 \ + --max-concurrency 1 \ + --ref-audio https://example.com/reference.wav \ + --ref-text "The exact transcript spoken in the reference audio." \ + --result-dir /tmp/voxcpm_clone_bench +``` + +## Notes + +- The benchmark uses `stream=true` and `response_format=pcm` so TTFP is measured from the first audio packet. +- `RTF < 1.0` means the server generates audio faster than real time. +- For `voxcpm_async_chunk.yaml`, keep concurrency at `1`. This matches native VoxCPM streaming more closely. +- Do not benchmark concurrent online streaming on `voxcpm_async_chunk.yaml`; use `voxcpm.yaml` for multi-request throughput runs. +- For the offline matrix mode, `--ref-audio` and `--ref-text` are required because clone cases are part of the fixed coverage set. diff --git a/benchmarks/voxcpm/vllm_omni/bench_tts_offline.py b/benchmarks/voxcpm/vllm_omni/bench_tts_offline.py new file mode 100644 index 000000000..a3bad3e69 --- /dev/null +++ b/benchmarks/voxcpm/vllm_omni/bench_tts_offline.py @@ -0,0 +1,890 @@ +"""Offline VoxCPM benchmark for vLLM Omni. + +Supports both: +- sync one-shot (Omni.generate) +- streaming (AsyncOmni.generate with async_chunk config) +- text-only synthesis +- voice cloning +- text/clone batch inputs from txt or jsonl +""" + +from __future__ import annotations + +import asyncio +import json +import logging +import os +import tempfile +import time +import uuid +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import torch +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import AsyncOmni, Omni + +REPO_ROOT = Path(__file__).resolve().parents[3] +DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml" +DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml" + +logger = logging.getLogger(__name__) + + +@dataclass(frozen=True, slots=True) +class PromptSpec: + text: str + label: str + ref_audio: str | None = None + ref_text: str | None = None + + +def _require_soundfile(): + try: + import soundfile as sf # type: ignore + except ModuleNotFoundError as exc: + raise RuntimeError( + "soundfile is required to write VoxCPM benchmark WAV outputs. Install it with: pip install soundfile" + ) from exc + return sf + + +def _build_prompt( + args, + *, + text: str, + ref_audio: str | None = None, + ref_text: str | None = None, + global_request_id: str | None = None, +) -> dict[str, Any]: + additional_information: dict[str, list[Any]] = { + "text": [text], + "cfg_value": [args.cfg_value], + "inference_timesteps": [args.inference_timesteps], + "min_len": [args.min_len], + "max_new_tokens": [args.max_new_tokens], + } + if args.streaming_prefix_len is not None: + additional_information["streaming_prefix_len"] = [args.streaming_prefix_len] + + if ref_audio: + additional_information["ref_audio"] = [ref_audio] + if ref_text: + additional_information["ref_text"] = [ref_text] + if global_request_id is not None: + additional_information["global_request_id"] = [global_request_id] + + return { + "prompt_token_ids": [1], + "additional_information": additional_information, + } + + +def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor: + audio = mm.get("audio", mm.get("model_outputs")) + if audio is None: + raise ValueError("No audio output found in multimodal output.") + if isinstance(audio, list): + parts = [torch.as_tensor(a).float().cpu().reshape(-1) for a in audio] + audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0) + if not isinstance(audio, torch.Tensor): + audio = torch.as_tensor(audio) + return audio.float().cpu().reshape(-1) + + +def _extract_sample_rate(mm: dict[str, Any]) -> int: + sr_raw = mm.get("sr", 24000) + if isinstance(sr_raw, list) and sr_raw: + sr_raw = sr_raw[-1] + if hasattr(sr_raw, "item"): + return int(sr_raw.item()) + return int(sr_raw) + + +def _emit_offline_metrics( + *, + request_id: str, + elapsed_s: float, + first_audio_elapsed: float | None, + audio_duration_s: float, +) -> None: + metrics = { + "request_id": request_id, + "ttfp_ms": round(first_audio_elapsed * 1000.0, 3) if first_audio_elapsed is not None else None, + "audio_duration_s": round(audio_duration_s, 6), + "rtf": round(elapsed_s / audio_duration_s, 6) if audio_duration_s > 0 else None, + } + print(f"[OfflineMetrics] {metrics}") + + +def _write_audio_tensor(output_path: Path, audio_tensor: Any, sample_rate: int) -> None: + sf = _require_soundfile() + if isinstance(audio_tensor, torch.Tensor): + audio_np = audio_tensor.float().cpu().clamp(-1.0, 1.0).numpy() + else: + audio_np = torch.as_tensor(audio_tensor).float().cpu().clamp(-1.0, 1.0).numpy() + sf.write( + output_path, + audio_np, + sample_rate, + format="WAV", + subtype="PCM_16", + ) + + +def _save_wav(mm: dict[str, Any], output_dir: Path, request_id: str) -> Path: + output_dir.mkdir(parents=True, exist_ok=True) + output_path = output_dir / f"output_{request_id}.wav" + _write_audio_tensor(output_path, _extract_audio_tensor(mm), _extract_sample_rate(mm)) + return output_path + + +def _iter_request_multimodal_outputs(request_output: Any): + outputs = getattr(request_output, "outputs", None) + if outputs: + for output in outputs: + mm = getattr(output, "multimodal_output", None) + if isinstance(mm, dict): + yield mm + + mm = getattr(request_output, "multimodal_output", None) + if isinstance(mm, dict): + yield mm + + +def _read_non_empty_lines(path: str) -> list[str]: + with open(path, encoding="utf-8") as f: + return [line.strip() for line in f if line.strip()] + + +def _load_prompt_specs(args) -> list[PromptSpec]: + specs: list[PromptSpec] = [] + + if args.txt_prompts is not None: + texts = _read_non_empty_lines(args.txt_prompts) + if not texts: + raise ValueError(f"No prompts found in {args.txt_prompts}") + for idx, text in enumerate(texts, start=1): + specs.append( + PromptSpec( + text=text, + label=f"item{idx:03d}", + ref_audio=args.ref_audio, + ref_text=args.ref_text, + ) + ) + return specs + + if args.jsonl_prompts is not None: + with open(args.jsonl_prompts, encoding="utf-8") as f: + for line_no, raw_line in enumerate(f, start=1): + line = raw_line.strip() + if not line: + continue + try: + item = json.loads(line) + except json.JSONDecodeError as exc: + raise ValueError(f"{args.jsonl_prompts}:{line_no} is not valid JSON: {exc}") from exc + if not isinstance(item, dict): + raise ValueError(f"{args.jsonl_prompts}:{line_no} must be a JSON object") + + text = item.get("text") + if not isinstance(text, str) or not text.strip(): + raise ValueError(f"{args.jsonl_prompts}:{line_no} requires non-empty string field 'text'") + + ref_audio = item.get("ref_audio", args.ref_audio) + ref_text = item.get("ref_text", args.ref_text) + if (ref_audio is None) != (ref_text is None): + raise ValueError( + f"{args.jsonl_prompts}:{line_no} must provide both 'ref_audio' and 'ref_text' together" + ) + + specs.append( + PromptSpec( + text=text.strip(), + label=f"item{len(specs) + 1:03d}", + ref_audio=ref_audio, + ref_text=ref_text, + ) + ) + + if not specs: + raise ValueError(f"No prompts found in {args.jsonl_prompts}") + return specs + + specs.append( + PromptSpec( + text=args.text, + label="item001", + ref_audio=args.ref_audio, + ref_text=args.ref_text, + ) + ) + return specs + + +def _build_prompt_for_spec(args, spec: PromptSpec, *, global_request_id: str | None = None) -> dict[str, Any]: + return _build_prompt( + args, + text=spec.text, + ref_audio=spec.ref_audio, + ref_text=spec.ref_text, + global_request_id=global_request_id, + ) + + +def _count_voice_clone_prompts(prompt_specs: list[PromptSpec]) -> int: + return sum(1 for spec in prompt_specs if spec.ref_audio is not None) + + +def _get_warmup_specs(prompt_specs: list[PromptSpec]) -> list[PromptSpec]: + return prompt_specs[:1] + + +def _extract_stream_finished(stage_output: Any) -> bool: + request_output = getattr(stage_output, "request_output", None) + request_finished = getattr(request_output, "finished", None) + if request_finished is not None: + return bool(request_finished) + return bool(getattr(stage_output, "finished", False)) + + +def _build_profiled_stage_config( + stage_configs_path: str, + profiler_dir: str, +) -> str: + stage_config_path = Path(stage_configs_path) + yaml_text = stage_config_path.read_text(encoding="utf-8") + injected_lines: list[str] = [] + injected_count = 0 + + for line in yaml_text.splitlines(): + injected_lines.append(line) + if line.strip() != "engine_args:": + continue + indent = line[: len(line) - len(line.lstrip())] + child_indent = indent + " " + grandchild_indent = child_indent + " " + injected_lines.extend( + [ + f"{child_indent}profiler_config:", + f'{grandchild_indent}profiler: "torch"', + f'{grandchild_indent}torch_profiler_dir: "{profiler_dir}"', + f"{grandchild_indent}torch_profiler_with_stack: true", + ] + ) + injected_count += 1 + + if injected_count == 0: + raise ValueError(f"No engine_args block found in stage config: {stage_configs_path}") + + tmp = tempfile.NamedTemporaryFile( + mode="w", + encoding="utf-8", + delete=False, + suffix=".yaml", + prefix=f"{stage_config_path.stem}_profile_", + ) + tmp.write("\n".join(injected_lines) + "\n") + tmp.close() + return tmp.name + + +def parse_args(): + parser = FlexibleArgumentParser( + description="Offline split-stage VoxCPM inference with vLLM Omni (auto sync/streaming by stage config)" + ) + parser.add_argument( + "--model", + type=str, + default=os.environ.get("VOXCPM_MODEL"), + help="Local VoxCPM model directory. Defaults to $VOXCPM_MODEL.", + ) + parser.add_argument( + "--text", + type=str, + default="This is a split-stage VoxCPM synthesis example running on vLLM Omni.", + help="Text to synthesize. Ignored when --txt-prompts or --jsonl-prompts is used.", + ) + parser.add_argument( + "--txt-prompts", + type=str, + default=None, + help="Path to a .txt file with one synthesis text per line.", + ) + parser.add_argument( + "--jsonl-prompts", + type=str, + default=None, + help=( + "Path to a .jsonl file. Each line must contain at least {'text': ...}; " + "clone rows can also set ref_audio/ref_text, and ref_text must be the " + "real transcript of ref_audio." + ), + ) + parser.add_argument( + "--ref-audio", + type=str, + default=None, + help=( + "Optional reference audio path for voice cloning. With --txt-prompts, " + "the same reference is applied to every line." + ), + ) + parser.add_argument( + "--ref-text", + type=str, + default=None, + help=( + "Real transcript of the reference audio. Placeholder text or mismatched " + "text will usually produce noisy/electronic clone audio." + ), + ) + parser.add_argument( + "--stage-configs-path", + type=str, + default=str(DEFAULT_STAGE_SYNC), + help="Stage config YAML path. Routing is selected only from this path.", + ) + parser.add_argument( + "--cfg-value", + type=float, + default=2.0, + help="Classifier-free guidance value for VoxCPM.", + ) + parser.add_argument( + "--inference-timesteps", + type=int, + default=10, + help="Number of inference timesteps.", + ) + parser.add_argument( + "--min-len", + type=int, + default=2, + help="Minimum generated token length.", + ) + parser.add_argument( + "--max-new-tokens", + type=int, + default=4096, + help="Maximum generated token length.", + ) + parser.add_argument( + "--streaming-prefix-len", + type=int, + default=None, + help="VoxCPM streaming window (optional, streaming mode only).", + ) + parser.add_argument( + "--output-dir", + type=str, + default=None, + help="Directory for output WAV files.", + ) + parser.add_argument( + "--stage-init-timeout", + type=int, + default=600, + help="Stage initialization timeout in seconds.", + ) + parser.add_argument( + "--log-stats", + dest="log_stats", + action="store_true", + help="Enable vLLM Omni stats logging.", + ) + parser.add_argument( + "--no-log-stats", + dest="log_stats", + action="store_false", + help="Disable vLLM Omni stats logging.", + ) + parser.set_defaults(log_stats=True) + parser.add_argument( + "--num-runs", + type=int, + default=1, + help="Number of full inference runs (same prompt each time). Default 1.", + ) + parser.add_argument( + "--warmup-runs", + type=int, + default=0, + help=( + "Optional number of warmup passes before measured runs. Warmup uses only " + "the first prompt and does not save outputs." + ), + ) + parser.add_argument( + "--enable-profiler", + action="store_true", + help=( + "Enable torch profiler for the configured stages. A temporary profiled " + "stage config is generated automatically." + ), + ) + parser.add_argument( + "--profiler-dir", + type=str, + default=None, + help="Directory for profiler traces. Defaults to /profiler when profiling is enabled.", + ) + parser.add_argument( + "--profiler-stages", + type=int, + nargs="*", + default=None, + help="Optional stage ids to profile. Defaults to all stages that have profiler_config.", + ) + parser.add_argument( + "--profiler-wait-seconds", + type=float, + default=30.0, + help="Seconds to wait after stop_profile for trace files to flush.", + ) + args = parser.parse_args() + + if not args.model: + parser.error("--model is required unless $VOXCPM_MODEL is set") + if args.txt_prompts is not None and args.jsonl_prompts is not None: + parser.error("--txt-prompts and --jsonl-prompts are mutually exclusive") + if (args.ref_audio is None) != (args.ref_text is None): + parser.error("--ref-audio and --ref-text must be provided together") + if args.num_runs < 1: + parser.error("--num-runs must be >= 1") + if args.warmup_runs < 0: + parser.error("--warmup-runs must be >= 0") + if args.output_dir is None: + args.output_dir = ( + "output_audio_streaming" if _is_streaming_stage_config(args.stage_configs_path) else "output_audio" + ) + if args.enable_profiler and args.profiler_dir is None: + args.profiler_dir = str(Path(args.output_dir) / "profiler") + try: + args.prompt_specs = _load_prompt_specs(args) + except ValueError as exc: + parser.error(str(exc)) + + return args + + +def _is_streaming_stage_config(stage_configs_path: str) -> bool: + cfg_name = Path(stage_configs_path).name.lower() + # Keep routing purely config-path based: + # - voxcpm.yaml => sync + # - voxcpm_async_chunk.yaml => streaming + return "async_chunk" in cfg_name + + +async def _collect_streaming_audio( + omni: AsyncOmni, + args: Any, + spec: PromptSpec, + request_id: str, + *, + phase_label: str, + prompt_index: int, + prompt_count: int, + print_prompt: bool = False, +) -> tuple[torch.Tensor, int, float, float | None]: + prompt = _build_prompt_for_spec(args, spec, global_request_id=request_id) + delta_chunks: list[torch.Tensor] = [] + sample_rate = 24000 + chunk_i = 0 + prev_total_samples = 0 + t_start = time.perf_counter() + first_audio_elapsed: float | None = None + + if print_prompt: + print(f"---prompt---:{prompt}") + + async for stage_output in omni.generate(prompt, request_id=request_id): + mm = getattr(stage_output, "multimodal_output", None) + if not isinstance(mm, dict): + ro = getattr(stage_output, "request_output", None) + if ro is None: + continue + mm = getattr(ro, "multimodal_output", None) + if not isinstance(mm, dict) and getattr(ro, "outputs", None): + seq = ro.outputs[0] + mm = getattr(seq, "multimodal_output", None) + if not isinstance(mm, dict): + continue + sample_rate = _extract_sample_rate(mm) + try: + w = _extract_audio_tensor(mm) + n = int(w.numel()) + if n == 0: + continue + finished = _extract_stream_finished(stage_output) + if n > prev_total_samples: + delta = w.reshape(-1)[prev_total_samples:] + prev_total_samples = n + elif finished and n == prev_total_samples: + delta = w.reshape(-1)[:0] + else: + delta = w.reshape(-1) + prev_total_samples += int(delta.numel()) + if int(delta.numel()) > 0: + delta_chunks.append(delta) + if first_audio_elapsed is None and int(delta.numel()) > 0: + first_audio_elapsed = time.perf_counter() - t_start + logger.info( + "%s prompt=%d/%d chunk=%d delta_samples=%d buf_len=%d finished=%s", + phase_label, + prompt_index + 1, + prompt_count, + chunk_i, + int(delta.numel()), + n, + finished, + ) + chunk_i += 1 + except ValueError: + if not _extract_stream_finished(stage_output): + logger.debug("skip non-audio partial output chunk=%d", chunk_i) + + if not delta_chunks: + raise RuntimeError("No audio chunks received; check stage config and logs.") + + audio_cat = torch.cat([c.reshape(-1) for c in delta_chunks], dim=0) + elapsed = time.perf_counter() - t_start + return audio_cat, sample_rate, elapsed, first_audio_elapsed + + +async def _abort_streaming_residual_work( + omni: AsyncOmni, + request_id: str, + *, + settle_seconds: float = 0.1, +) -> None: + """Stop any late stage-0 work once the final audio has been collected.""" + await omni.engine.abort_async([request_id]) + if settle_seconds > 0: + await asyncio.sleep(settle_seconds) + + +async def _run_streaming_single( + omni: AsyncOmni, + args: Any, + spec: PromptSpec, + output_dir: Path, + request_id: str, + *, + run_index: int, + num_runs: int, + prompt_index: int, + prompt_count: int, +) -> Path: + audio_cat, sample_rate, elapsed, first_audio_elapsed = await _collect_streaming_audio( + omni, + args, + spec, + request_id, + phase_label=f"run={run_index + 1}/{num_runs}", + prompt_index=prompt_index, + prompt_count=prompt_count, + print_prompt=(run_index == 0 and prompt_index == 0), + ) + await _abort_streaming_residual_work(omni, request_id) + output_path = output_dir / f"output_run{run_index + 1}_{spec.label}.wav" + _write_audio_tensor(output_path, audio_cat, sample_rate) + audio_duration_s = float(audio_cat.numel()) / float(sample_rate) if sample_rate > 0 else 0.0 + ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else "" + rtf_text = f", rtf={elapsed / audio_duration_s:.3f}" if audio_duration_s > 0 else "" + print( + f"Saved (streaming) run {run_index + 1}/{num_runs}, " + f"prompt {prompt_index + 1}/{prompt_count}: {output_path} ({elapsed:.2f}s{ttfp_text}{rtf_text})" + ) + _emit_offline_metrics( + request_id=request_id, + elapsed_s=elapsed, + first_audio_elapsed=first_audio_elapsed, + audio_duration_s=audio_duration_s, + ) + return output_path + + +async def _run_streaming_warmup(args, omni: AsyncOmni) -> None: + if args.warmup_runs == 0: + return + + warmup_specs = _get_warmup_specs(args.prompt_specs) + print( + f"Warmup: {args.warmup_runs} run(s) using the first prompt " + f"({len(warmup_specs)} prompt(s)); outputs will be discarded." + ) + for warmup_index in range(args.warmup_runs): + t_warmup = time.perf_counter() + tasks = [] + request_ids: list[str] = [] + for prompt_index, spec in enumerate(warmup_specs): + request_id = f"warmup_stream_{warmup_index + 1}_{spec.label}_{uuid.uuid4().hex[:8]}" + request_ids.append(request_id) + tasks.append( + _collect_streaming_audio( + omni, + args, + spec, + request_id, + phase_label=f"warmup={warmup_index + 1}/{args.warmup_runs}", + prompt_index=prompt_index, + prompt_count=len(warmup_specs), + ) + ) + results = await asyncio.gather(*tasks) + for request_id in request_ids: + await _abort_streaming_residual_work(omni, request_id) + total_samples = sum(int(audio.numel()) for audio, _, _, _ in results) + warmup_ttfps = [ttfp for _, _, _, ttfp in results if ttfp is not None] + ttfp_text = f", ttfp={min(warmup_ttfps):.2f}s" if warmup_ttfps else "" + print( + f"Warmup (streaming) {warmup_index + 1}/{args.warmup_runs} finished: " + f"{len(results)} prompt(s), {total_samples} sample(s) " + f"({time.perf_counter() - t_warmup:.2f}s{ttfp_text})" + ) + + +async def _run_streaming(args) -> list[Path]: + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + + omni = AsyncOmni( + model=args.model, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + + await _run_streaming_warmup(args, omni) + profiler_started = False + if args.enable_profiler: + profile_prefix = f"voxcpm_streaming_{int(time.time())}" + stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured" + print(f"Starting profiler (streaming): stages={stages_text}, dir={args.profiler_dir}") + await omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages) + profiler_started = True + t_total = time.perf_counter() + total_elapsed = 0.0 + paths: list[Path] = [] + prompt_specs: list[PromptSpec] = args.prompt_specs + try: + for run in range(args.num_runs): + for prompt_index, spec in enumerate(prompt_specs): + request_id = f"stream_{run + 1}_{spec.label}_{uuid.uuid4().hex[:8]}" + paths.append( + await _run_streaming_single( + omni, + args, + spec, + output_dir, + request_id, + run_index=run, + num_runs=args.num_runs, + prompt_index=prompt_index, + prompt_count=len(prompt_specs), + ) + ) + total_elapsed = time.perf_counter() - t_total + finally: + if profiler_started: + print("Stopping profiler (streaming)...") + await omni.stop_profile(stages=args.profiler_stages) + if args.profiler_wait_seconds > 0: + print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...") + await asyncio.sleep(args.profiler_wait_seconds) + + print( + f"All streaming runs finished: {args.num_runs} run(s), " + f"{len(prompt_specs)} prompt(s), {len(paths)} file(s) in {total_elapsed:.2f}s total" + ) + return paths + + +def _run_sync(args) -> list[Path]: + output_dir = Path(args.output_dir) + + omni = Omni( + model=args.model, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + + def _run_sync_single( + spec: PromptSpec, + *, + request_prefix: str, + save_outputs: bool, + run_index: int | None = None, + ) -> tuple[list[Path], int, float | None, float, float, str]: + global_request_id = f"{request_prefix}_{spec.label}" + prompt = _build_prompt_for_spec(args, spec, global_request_id=global_request_id) + if save_outputs and run_index == 0 and spec.label == "item001": + print(f"---prompt---:{prompt}") + + saved_paths: list[Path] = [] + output_count = 0 + first_audio_elapsed: float | None = None + total_audio_duration_s = 0.0 + metrics_request_id = global_request_id + t_start = time.perf_counter() + for stage_outputs in omni.generate(prompt): + request_output = stage_outputs.request_output + if request_output is None: + continue + request_output_id = getattr(request_output, "request_id", None) + if isinstance(request_output_id, str) and request_output_id: + metrics_request_id = request_output_id + for j, mm in enumerate(_iter_request_multimodal_outputs(request_output)): + output_count += 1 + if first_audio_elapsed is None: + try: + audio_tensor = _extract_audio_tensor(mm) + if int(audio_tensor.numel()) > 0: + first_audio_elapsed = time.perf_counter() - t_start + total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm)) + except ValueError: + pass + else: + try: + audio_tensor = _extract_audio_tensor(mm) + total_audio_duration_s += float(audio_tensor.numel()) / float(_extract_sample_rate(mm)) + except ValueError: + pass + if not save_outputs: + continue + save_stem = f"run{run_index + 1}_{spec.label}" if j == 0 else f"run{run_index + 1}_{spec.label}_{j}" + saved_paths.append(_save_wav(mm, output_dir, save_stem)) + + if output_count == 0: + raise RuntimeError("No output from Omni.generate") + elapsed_s = time.perf_counter() - t_start + return saved_paths, output_count, first_audio_elapsed, elapsed_s, total_audio_duration_s, metrics_request_id + + if args.warmup_runs: + warmup_specs = _get_warmup_specs(args.prompt_specs) + print( + f"Warmup: {args.warmup_runs} run(s) using the first prompt " + f"({len(warmup_specs)} prompt(s)); outputs will be discarded." + ) + for warmup_index in range(args.warmup_runs): + t_warmup = time.perf_counter() + _, output_count, first_audio_elapsed, elapsed_s, audio_duration_s, _ = _run_sync_single( + warmup_specs[0], + request_prefix=f"warmup_sync{warmup_index + 1}", + save_outputs=False, + ) + ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else "" + rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else "" + print( + f"Warmup (sync) {warmup_index + 1}/{args.warmup_runs} finished: " + f"{output_count} output(s) ({time.perf_counter() - t_warmup:.2f}s{ttfp_text}{rtf_text})" + ) + + profiler_started = False + if args.enable_profiler: + profile_prefix = f"voxcpm_sync_{int(time.time())}" + stages_text = args.profiler_stages if args.profiler_stages is not None else "all-configured" + print(f"Starting profiler (sync): stages={stages_text}, dir={args.profiler_dir}") + omni.start_profile(profile_prefix=profile_prefix, stages=args.profiler_stages) + profiler_started = True + + t_total = time.perf_counter() + total_elapsed = 0.0 + saved_paths: list[Path] = [] + prompt_specs: list[PromptSpec] = args.prompt_specs + try: + for run in range(args.num_runs): + t_run = time.perf_counter() + run_paths: list[Path] = [] + for prompt_index, spec in enumerate(prompt_specs): + prompt_paths, _, first_audio_elapsed, elapsed_s, audio_duration_s, metrics_request_id = ( + _run_sync_single( + spec, + request_prefix=f"sync_run{run + 1}_{prompt_index + 1:03d}", + save_outputs=True, + run_index=run, + ) + ) + run_paths.extend(prompt_paths) + ttfp_text = f", ttfp={first_audio_elapsed:.2f}s" if first_audio_elapsed is not None else "" + rtf_text = f", rtf={elapsed_s / audio_duration_s:.3f}" if audio_duration_s > 0 else "" + print( + f"Saved (sync) run {run + 1}/{args.num_runs}, " + f"prompt {prompt_index + 1}/{len(prompt_specs)}: {len(prompt_paths)} file(s){ttfp_text}{rtf_text}" + ) + _emit_offline_metrics( + request_id=metrics_request_id, + elapsed_s=elapsed_s, + first_audio_elapsed=first_audio_elapsed, + audio_duration_s=audio_duration_s, + ) + + saved_paths.extend(run_paths) + print( + f"Run {run + 1}/{args.num_runs} finished: {len(run_paths)} file(s) ({time.perf_counter() - t_run:.2f}s)" + ) + for path in run_paths: + print(f" {path}") + + total_elapsed = time.perf_counter() - t_total + finally: + if profiler_started: + print("Stopping profiler (sync)...") + omni.stop_profile(stages=args.profiler_stages) + if args.profiler_wait_seconds > 0: + print(f"Waiting {args.profiler_wait_seconds:.1f}s for profiler traces to flush...") + time.sleep(args.profiler_wait_seconds) + + print( + f"All sync runs finished: {args.num_runs} run(s), " + f"{len(prompt_specs)} prompt(s), {len(saved_paths)} file(s) in {total_elapsed:.2f}s total" + ) + return saved_paths + + +def main(args) -> int: + logging.basicConfig(level=logging.INFO) + profiled_stage_config_path: str | None = None + original_stage_config_path = args.stage_configs_path + if args.enable_profiler: + Path(args.profiler_dir).mkdir(parents=True, exist_ok=True) + profiled_stage_config_path = _build_profiled_stage_config( + args.stage_configs_path, + str(Path(args.profiler_dir).resolve()), + ) + args.stage_configs_path = profiled_stage_config_path + + is_streaming = _is_streaming_stage_config(args.stage_configs_path) + voice_clone_count = _count_voice_clone_prompts(args.prompt_specs) + print(f"Model: {args.model}") + print(f"Stage config: {original_stage_config_path}") + print(f"Route: {'streaming' if is_streaming else 'sync'} (from stage-configs-path)") + print(f"Prompt count: {len(args.prompt_specs)}") + print("Batch mode: sequential (aligned with native VoxCPM)") + print(f"Warmup runs: {args.warmup_runs}") + print(f"Voice cloning prompts: {voice_clone_count}/{len(args.prompt_specs)}") + if args.enable_profiler: + print(f"Profiler: enabled (dir={args.profiler_dir}, stages={args.profiler_stages or 'all-configured'})") + print(f"Profiled stage config: {args.stage_configs_path}") + if voice_clone_count: + print("Voice cloning note: --ref-text/ref_text must match the spoken content of the reference audio.") + print(f"Num runs: {args.num_runs}") + try: + if is_streaming: + asyncio.run(_run_streaming(args)) + else: + _run_sync(args) + finally: + if profiled_stage_config_path is not None and os.path.exists(profiled_stage_config_path): + os.unlink(profiled_stage_config_path) + return 0 + + +if __name__ == "__main__": + os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + raise SystemExit(main(parse_args())) diff --git a/benchmarks/voxcpm/vllm_omni/bench_tts_serve.py b/benchmarks/voxcpm/vllm_omni/bench_tts_serve.py new file mode 100644 index 000000000..816df3279 --- /dev/null +++ b/benchmarks/voxcpm/vllm_omni/bench_tts_serve.py @@ -0,0 +1,283 @@ +"""Benchmark VoxCPM via /v1/audio/speech. + +Reports TTFP (time to first packet), E2E latency, and RTF (real-time factor). +""" + +from __future__ import annotations + +import argparse +import asyncio +import json +import time +from dataclasses import asdict, dataclass, field +from datetime import datetime +from pathlib import Path + +import aiohttp +import numpy as np +from tqdm.asyncio import tqdm + +DEFAULT_MODEL = "OpenBMB/VoxCPM1.5" +DEFAULT_SAMPLE_RATE = 24000 +PROMPTS = [ + "Hello, welcome to the VoxCPM speech benchmark.", + "This is a short benchmark prompt for online text-to-speech generation.", + "The quick brown fox jumps over the lazy dog near the riverbank.", + "Please remember to bring your identification documents tomorrow morning.", + "Learning a new language takes patience, practice, and curiosity.", + "This benchmark reports TTFP and RTF for the VoxCPM online serving path.", +] + + +@dataclass +class RequestResult: + success: bool = False + ttfp: float = 0.0 + e2e: float = 0.0 + audio_bytes: int = 0 + audio_duration: float = 0.0 + rtf: float = 0.0 + prompt: str = "" + error: str = "" + + +@dataclass +class BenchmarkResult: + concurrency: int = 0 + num_prompts: int = 0 + completed: int = 0 + failed: int = 0 + duration_s: float = 0.0 + mean_ttfp_ms: float = 0.0 + median_ttfp_ms: float = 0.0 + p95_ttfp_ms: float = 0.0 + mean_e2e_ms: float = 0.0 + median_e2e_ms: float = 0.0 + p95_e2e_ms: float = 0.0 + mean_rtf: float = 0.0 + median_rtf: float = 0.0 + p95_rtf: float = 0.0 + total_audio_duration_s: float = 0.0 + request_throughput: float = 0.0 + per_request: list[dict[str, float | str]] = field(default_factory=list) + + +def pcm_bytes_to_duration(num_bytes: int, sample_rate: int = DEFAULT_SAMPLE_RATE, sample_width: int = 2) -> float: + num_samples = num_bytes / sample_width + return num_samples / sample_rate + + +async def send_tts_request( + session: aiohttp.ClientSession, + api_url: str, + *, + model: str, + prompt: str, + ref_audio: str | None, + ref_text: str | None, + pbar: tqdm | None = None, +) -> RequestResult: + payload: dict[str, object] = { + "model": model, + "input": prompt, + "stream": True, + "response_format": "pcm", + } + if ref_audio is not None: + payload["ref_audio"] = ref_audio + if ref_text is not None: + payload["ref_text"] = ref_text + + result = RequestResult(prompt=prompt) + started_at = time.perf_counter() + + try: + async with session.post(api_url, json=payload) as response: + if response.status != 200: + result.error = f"HTTP {response.status}: {await response.text()}" + return result + + first_chunk = True + total_bytes = 0 + async for chunk in response.content.iter_any(): + if not chunk: + continue + if first_chunk: + result.ttfp = time.perf_counter() - started_at + first_chunk = False + total_bytes += len(chunk) + + result.e2e = time.perf_counter() - started_at + result.audio_bytes = total_bytes + result.audio_duration = pcm_bytes_to_duration(total_bytes) + if result.audio_duration > 0: + result.rtf = result.e2e / result.audio_duration + result.success = True + except Exception as e: + result.error = str(e) + result.e2e = time.perf_counter() - started_at + + if pbar is not None: + pbar.update(1) + return result + + +async def run_benchmark( + *, + host: str, + port: int, + model: str, + num_prompts: int, + max_concurrency: int, + num_warmups: int, + ref_audio: str | None, + ref_text: str | None, +) -> BenchmarkResult: + api_url = f"http://{host}:{port}/v1/audio/speech" + connector = aiohttp.TCPConnector(limit=max_concurrency, limit_per_host=max_concurrency, keepalive_timeout=60) + timeout = aiohttp.ClientTimeout(total=600) + + async with aiohttp.ClientSession(connector=connector, timeout=timeout) as session: + if num_warmups > 0: + print(f" Warming up with {num_warmups} requests...") + warmup_tasks = [ + send_tts_request( + session, + api_url, + model=model, + prompt=PROMPTS[i % len(PROMPTS)], + ref_audio=ref_audio, + ref_text=ref_text, + ) + for i in range(num_warmups) + ] + await asyncio.gather(*warmup_tasks) + print(" Warmup done.") + + request_prompts = [PROMPTS[i % len(PROMPTS)] for i in range(num_prompts)] + semaphore = asyncio.Semaphore(max_concurrency) + pbar = tqdm(total=num_prompts, desc=f" concurrency={max_concurrency}") + + async def limited_request(prompt: str) -> RequestResult: + async with semaphore: + return await send_tts_request( + session, + api_url, + model=model, + prompt=prompt, + ref_audio=ref_audio, + ref_text=ref_text, + pbar=pbar, + ) + + started_at = time.perf_counter() + results = await asyncio.gather(*[asyncio.create_task(limited_request(prompt)) for prompt in request_prompts]) + duration = time.perf_counter() - started_at + pbar.close() + + succeeded = [result for result in results if result.success] + bench = BenchmarkResult( + concurrency=max_concurrency, + num_prompts=num_prompts, + completed=len(succeeded), + failed=len(results) - len(succeeded), + duration_s=duration, + ) + + if not succeeded: + return bench + + ttfps = np.array([result.ttfp * 1000 for result in succeeded], dtype=np.float64) + e2es = np.array([result.e2e * 1000 for result in succeeded], dtype=np.float64) + rtfs = np.array([result.rtf for result in succeeded], dtype=np.float64) + audio_durations = np.array([result.audio_duration for result in succeeded], dtype=np.float64) + + bench.mean_ttfp_ms = float(np.mean(ttfps)) + bench.median_ttfp_ms = float(np.median(ttfps)) + bench.p95_ttfp_ms = float(np.percentile(ttfps, 95)) + bench.mean_e2e_ms = float(np.mean(e2es)) + bench.median_e2e_ms = float(np.median(e2es)) + bench.p95_e2e_ms = float(np.percentile(e2es, 95)) + bench.mean_rtf = float(np.mean(rtfs)) + bench.median_rtf = float(np.median(rtfs)) + bench.p95_rtf = float(np.percentile(rtfs, 95)) + bench.total_audio_duration_s = float(np.sum(audio_durations)) + bench.request_throughput = len(succeeded) / duration if duration > 0 else 0.0 + bench.per_request = [ + { + "prompt": result.prompt, + "ttfp_ms": result.ttfp * 1000, + "e2e_ms": result.e2e * 1000, + "rtf": result.rtf, + "audio_duration_s": result.audio_duration, + } + for result in succeeded + ] + + return bench + + +def print_summary(result: BenchmarkResult) -> None: + width = 54 + print("") + print("=" * width) + print(f"{'VoxCPM Serving Benchmark':^{width}}") + print("=" * width) + print(f"concurrency : {result.concurrency}") + print(f"requests : {result.completed}/{result.num_prompts} succeeded") + print(f"wall time (s) : {result.duration_s:.3f}") + print(f"mean TTFP (ms) : {result.mean_ttfp_ms:.2f}") + print(f"p95 TTFP (ms) : {result.p95_ttfp_ms:.2f}") + print(f"mean E2E (ms) : {result.mean_e2e_ms:.2f}") + print(f"p95 E2E (ms) : {result.p95_e2e_ms:.2f}") + print(f"mean RTF : {result.mean_rtf:.3f}") + print(f"p95 RTF : {result.p95_rtf:.3f}") + print(f"request throughput : {result.request_throughput:.2f} req/s") + print("=" * width) + + +async def main_async(args) -> None: + result_dir = Path(args.result_dir) + result_dir.mkdir(parents=True, exist_ok=True) + + all_results: list[BenchmarkResult] = [] + for concurrency in args.max_concurrency: + result = await run_benchmark( + host=args.host, + port=args.port, + model=args.model, + num_prompts=args.num_prompts, + max_concurrency=concurrency, + num_warmups=args.num_warmups, + ref_audio=args.ref_audio, + ref_text=args.ref_text, + ) + print_summary(result) + all_results.append(result) + + payload = { + "model": args.model, + "created_at": datetime.utcnow().isoformat() + "Z", + "results": [asdict(result) for result in all_results], + } + result_path = result_dir / "bench_tts_serve.json" + result_path.write_text(json.dumps(payload, indent=2), encoding="utf-8") + print(f"Saved results to: {result_path}") + + +def parse_args(): + parser = argparse.ArgumentParser(description="Benchmark VoxCPM via /v1/audio/speech") + parser.add_argument("--host", default="127.0.0.1", help="Server host") + parser.add_argument("--port", type=int, default=8091, help="Server port") + parser.add_argument("--model", default=DEFAULT_MODEL, help="Model name or path") + parser.add_argument("--num-prompts", type=int, default=20, help="Number of prompts to send") + parser.add_argument("--max-concurrency", type=int, nargs="+", default=[1], help="Concurrency levels to benchmark") + parser.add_argument("--num-warmups", type=int, default=3, help="Warmup request count") + parser.add_argument("--ref-audio", default=None, help="Reference audio URL or data URL for voice cloning") + parser.add_argument("--ref-text", default=None, help="Reference audio transcript for voice cloning") + parser.add_argument("--result-dir", default="results", help="Directory to save benchmark JSON") + return parser.parse_args() + + +if __name__ == "__main__": + asyncio.run(main_async(parse_args())) diff --git a/benchmarks/voxcpm/vllm_omni/run_offline_matrix.py b/benchmarks/voxcpm/vllm_omni/run_offline_matrix.py new file mode 100644 index 000000000..cee46c0f8 --- /dev/null +++ b/benchmarks/voxcpm/vllm_omni/run_offline_matrix.py @@ -0,0 +1,303 @@ +"""Run the full offline VoxCPM smoke matrix. + +This script keeps the old `test.py` coverage, but delegates each case to +`bench_tts_offline.py` so the benchmark runner itself stays focused on a +single execution path. +""" + +from __future__ import annotations + +import shlex +import subprocess +import sys +import time +from dataclasses import dataclass +from pathlib import Path + +from vllm.utils.argparse_utils import FlexibleArgumentParser + +REPO_ROOT = Path(__file__).resolve().parents[3] +BENCH_SCRIPT = Path(__file__).with_name("bench_tts_offline.py") +DEFAULT_STAGE_ASYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm_async_chunk.yaml" +DEFAULT_STAGE_SYNC = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml" +DEFAULT_OUTPUT_ROOT = BENCH_SCRIPT.parents[1] / "results" / "offline_matrix" + +SINGLE_TTS_TEXT = "This is a single text-to-speech smoke test for VoxCPM on vLLM Omni." +SINGLE_CLONE_TEXT = "This sentence is synthesized with the cloned voice for validation." +BATCH_TTS_TEXTS = [ + "The first batch text-to-speech sample validates sequential batch execution.", + "The second batch text-to-speech sample checks another prompt in the same file.", + "The third batch text-to-speech sample completes the sequential batch path.", +] +BATCH_CLONE_TEXTS = [ + "The first cloned sample validates sequential batch voice cloning.", + "The second cloned sample checks the same reference voice on another prompt.", + "The third cloned sample finishes the shared-reference clone batch path.", +] + + +@dataclass(frozen=True, slots=True) +class ModeSpec: + name: str + stage_config: Path + + +@dataclass(frozen=True, slots=True) +class CaseSpec: + name: str + warmup_runs: int + prompt_kind: str + voice_clone: bool + + +@dataclass(frozen=True, slots=True) +class CaseResult: + mode: str + case: str + returncode: int + elapsed_s: float + output_dir: Path + log_path: Path + + @property + def ok(self) -> bool: + return self.returncode == 0 + + +MODE_SPECS = [ + ModeSpec(name="streaming", stage_config=DEFAULT_STAGE_ASYNC), + ModeSpec(name="sync", stage_config=DEFAULT_STAGE_SYNC), +] + +CASE_SPECS = [ + CaseSpec(name="warmup_single_tts", warmup_runs=1, prompt_kind="single", voice_clone=False), + CaseSpec(name="warmup_single_clone", warmup_runs=1, prompt_kind="single", voice_clone=True), + CaseSpec(name="warmup_batch_tts", warmup_runs=1, prompt_kind="batch", voice_clone=False), + CaseSpec(name="warmup_batch_clone", warmup_runs=1, prompt_kind="batch", voice_clone=True), + CaseSpec(name="cold_single_tts", warmup_runs=0, prompt_kind="single", voice_clone=False), + CaseSpec(name="cold_single_clone", warmup_runs=0, prompt_kind="single", voice_clone=True), +] + + +def _write_lines(path: Path, lines: list[str]) -> None: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text("\n".join(lines) + "\n", encoding="utf-8") + + +def _prepare_batch_inputs(output_root: Path) -> tuple[Path, Path]: + input_dir = output_root / "inputs" + batch_tts_path = input_dir / "batch_tts_prompts.txt" + batch_clone_path = input_dir / "batch_clone_prompts.txt" + _write_lines(batch_tts_path, BATCH_TTS_TEXTS) + _write_lines(batch_clone_path, BATCH_CLONE_TEXTS) + return batch_tts_path, batch_clone_path + + +def _base_command(args, mode: ModeSpec, output_dir: Path) -> list[str]: + cmd = [ + args.python, + str(BENCH_SCRIPT), + "--model", + args.model, + "--stage-configs-path", + str(mode.stage_config), + "--output-dir", + str(output_dir), + "--num-runs", + str(args.num_runs), + "--stage-init-timeout", + str(args.stage_init_timeout), + ] + cmd.append("--log-stats" if args.log_stats else "--no-log-stats") + cmd.extend(["--cfg-value", str(args.cfg_value)]) + cmd.extend(["--inference-timesteps", str(args.inference_timesteps)]) + cmd.extend(["--min-len", str(args.min_len)]) + cmd.extend(["--max-new-tokens", str(args.max_new_tokens)]) + if args.streaming_prefix_len is not None: + cmd.extend(["--streaming-prefix-len", str(args.streaming_prefix_len)]) + if args.enable_profiler: + profiler_dir = Path(args.profiler_dir) if args.profiler_dir is not None else (output_dir / "profiler") + cmd.append("--enable-profiler") + cmd.extend(["--profiler-dir", str(profiler_dir)]) + cmd.extend(["--profiler-wait-seconds", str(args.profiler_wait_seconds)]) + if args.profiler_stages is not None: + cmd.append("--profiler-stages") + cmd.extend(str(stage_id) for stage_id in args.profiler_stages) + return cmd + + +def _build_case_command( + args, + mode: ModeSpec, + case: CaseSpec, + *, + batch_tts_path: Path, + batch_clone_path: Path, + output_dir: Path, +) -> list[str]: + cmd = _base_command(args, mode, output_dir) + cmd.extend(["--warmup-runs", str(case.warmup_runs)]) + if case.prompt_kind == "single": + cmd.extend(["--text", SINGLE_CLONE_TEXT if case.voice_clone else SINGLE_TTS_TEXT]) + else: + cmd.extend(["--txt-prompts", str(batch_clone_path if case.voice_clone else batch_tts_path)]) + if case.voice_clone: + cmd.extend(["--ref-audio", args.ref_audio, "--ref-text", args.ref_text]) + return cmd + + +def _run_case( + args, + mode: ModeSpec, + case: CaseSpec, + *, + batch_tts_path: Path, + batch_clone_path: Path, + output_root: Path, +) -> CaseResult: + case_output_dir = output_root / mode.name / case.name + case_output_dir.mkdir(parents=True, exist_ok=True) + case_log_path = case_output_dir / "run.log" + cmd = _build_case_command( + args, + mode, + case, + batch_tts_path=batch_tts_path, + batch_clone_path=batch_clone_path, + output_dir=case_output_dir, + ) + + print() + print("=" * 80) + print(f"[{mode.name}] {case.name}") + print(f"Output directory: {case_output_dir}") + print(shlex.join(cmd)) + + start = time.perf_counter() + with case_log_path.open("w", encoding="utf-8") as log_fp: + process = subprocess.Popen( + cmd, + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + bufsize=1, + ) + assert process.stdout is not None + for line in process.stdout: + print(line, end="") + log_fp.write(line) + process.wait() + + elapsed_s = time.perf_counter() - start + status = "PASS" if (process.returncode or 0) == 0 else f"FAIL({process.returncode})" + print(f"[{mode.name}] {case.name} -> {status} ({elapsed_s:.2f}s)") + return CaseResult( + mode=mode.name, + case=case.name, + returncode=int(process.returncode or 0), + elapsed_s=elapsed_s, + output_dir=case_output_dir, + log_path=case_log_path, + ) + + +def parse_args(): + parser = FlexibleArgumentParser(description="Run the full offline VoxCPM smoke matrix.") + parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.") + parser.add_argument("--ref-audio", type=str, required=True, help="Reference audio path for clone cases.") + parser.add_argument("--ref-text", type=str, required=True, help="Exact transcript spoken in --ref-audio.") + parser.add_argument("--output-root", type=str, default=str(DEFAULT_OUTPUT_ROOT), help="Root directory for outputs.") + parser.add_argument("--python", type=str, default=sys.executable, help="Python executable used to launch cases.") + parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.") + parser.add_argument("--log-stats", dest="log_stats", action="store_true", help="Enable vLLM Omni stats logging.") + parser.add_argument( + "--no-log-stats", + dest="log_stats", + action="store_false", + help="Disable vLLM Omni stats logging.", + ) + parser.set_defaults(log_stats=True) + parser.add_argument("--num-runs", type=int, default=1, help="Number of measured runs per case.") + parser.add_argument("--cfg-value", type=float, default=2.0, help="Classifier-free guidance value for VoxCPM.") + parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of inference timesteps.") + parser.add_argument("--min-len", type=int, default=2, help="Minimum generated token length.") + parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum generated token length.") + parser.add_argument( + "--streaming-prefix-len", + type=int, + default=None, + help="Optional VoxCPM streaming window passed to streaming cases.", + ) + parser.add_argument("--enable-profiler", action="store_true", help="Enable torch profiler for each case.") + parser.add_argument( + "--profiler-dir", + type=str, + default=None, + help="Profiler output root. Defaults to /profiler.", + ) + parser.add_argument( + "--profiler-stages", + type=int, + nargs="*", + default=None, + help="Optional stage ids to profile. Defaults to all configured stages.", + ) + parser.add_argument( + "--profiler-wait-seconds", + type=float, + default=30.0, + help="Seconds to wait after stopping profiler for traces to flush.", + ) + args = parser.parse_args() + if args.num_runs < 1: + parser.error("--num-runs must be >= 1") + return args + + +def main(args) -> int: + output_root = Path(args.output_root) + output_root.mkdir(parents=True, exist_ok=True) + batch_tts_path, batch_clone_path = _prepare_batch_inputs(output_root) + + print(f"Model: {args.model}") + print(f"Reference audio: {args.ref_audio}") + print(f"Reference text: {args.ref_text}") + print(f"Python: {args.python}") + print(f"Output root: {output_root}") + print(f"Cases: {len(MODE_SPECS) * len(CASE_SPECS)}") + + results: list[CaseResult] = [] + for mode in MODE_SPECS: + for case in CASE_SPECS: + results.append( + _run_case( + args, + mode, + case, + batch_tts_path=batch_tts_path, + batch_clone_path=batch_clone_path, + output_root=output_root, + ) + ) + + failed = [result for result in results if not result.ok] + print() + print("=" * 80) + print("Summary:") + for result in results: + status = "PASS" if result.ok else f"FAIL({result.returncode})" + print(f"- [{result.mode}] {result.case}: {status} ({result.elapsed_s:.2f}s)") + print(f" output_dir={result.output_dir}") + print(f" log={result.log_path}") + + print(f"Passed: {len(results) - len(failed)}/{len(results)}") + if failed: + print("Failed cases:") + for result in failed: + print(f"- [{result.mode}] {result.case}: see {result.log_path}") + return 1 + return 0 + + +if __name__ == "__main__": + raise SystemExit(main(parse_args())) diff --git a/examples/offline_inference/voxcpm/README.md b/examples/offline_inference/voxcpm/README.md new file mode 100644 index 000000000..1eaea9b0d --- /dev/null +++ b/examples/offline_inference/voxcpm/README.md @@ -0,0 +1,123 @@ +# VoxCPM Offline Example + +This directory contains the minimal offline VoxCPM example for vLLM Omni. + +`end2end.py` is intentionally small and only covers: + +- single text-to-speech +- single voice cloning with `ref_audio` + `ref_text` +- non-streaming with `vllm_omni/model_executor/stage_configs/voxcpm.yaml` +- streaming with `vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml` + +Advanced workflows were moved out of the getting-started example: + +- `benchmarks/voxcpm/vllm_omni/bench_tts_offline.py`: warmup, batch prompts, profiler, offline TTFP / RTF +- `benchmarks/voxcpm/vllm_omni/run_offline_matrix.py`: fixed offline smoke matrix +- `benchmarks/voxcpm/`: benchmark scripts and benchmark docs + +## Prerequisites + +Install VoxCPM in one of these ways: + +```bash +pip install voxcpm +``` + +or point vLLM Omni to the local VoxCPM source tree: + +```bash +export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src +``` + +The example writes WAV files with `soundfile`: + +```bash +pip install soundfile +``` + +## Model Path + +Pass the native VoxCPM model directory directly: + +```bash +export VOXCPM_MODEL=/path/to/voxcpm-model +``` + +If the native VoxCPM `config.json` does not contain HuggingFace metadata such as +`model_type`, prepare a persistent HF-compatible config directory and point the +stage configs to it with `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH`: + +```bash +export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config +mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH" +cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json" +cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true +python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)' +``` + +If the model directory itself already has `model_type`, this extra directory is +not required. + +## Quick Start + +Single text-to-speech, non-streaming: + +```bash +python examples/offline_inference/voxcpm/end2end.py \ + --model "$VOXCPM_MODEL" \ + --text "This is a split-stage VoxCPM synthesis example running on vLLM Omni." +``` + +Single voice cloning, non-streaming: + +```bash +python examples/offline_inference/voxcpm/end2end.py \ + --model "$VOXCPM_MODEL" \ + --text "This sentence is synthesized with a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript spoken in reference.wav." +``` + +Streaming: + +```bash +python examples/offline_inference/voxcpm/end2end.py \ + --model "$VOXCPM_MODEL" \ + --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \ + --text "This is a split-stage VoxCPM streaming example running on vLLM Omni." +``` + +By default, `end2end.py` writes to `output_audio/` for non-streaming and +`output_audio_streaming/` for streaming. + +## Advanced Workflows + +Use `benchmarks/voxcpm/vllm_omni/bench_tts_offline.py` when you need: + +- warmup runs +- prompt files +- batch JSONL inputs +- profiler injection +- offline TTFP / RTF emission + +Use `benchmarks/voxcpm/vllm_omni/run_offline_matrix.py` when you need the fixed offline smoke matrix that previously lived in `test.py`. + +Full matrix benchmark example: + +```bash +python benchmarks/voxcpm/vllm_omni/run_offline_matrix.py \ + --model "$VOXCPM_MODEL" \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript spoken in reference.wav." +``` + +For online serving examples, see [examples/online_serving/voxcpm](../../online_serving/voxcpm/README.md). + +For benchmark reporting, see [benchmarks/voxcpm](../../../benchmarks/voxcpm/README.md). + +## Notes + +- `voxcpm.yaml` is the default non-streaming stage config. +- `voxcpm_async_chunk.yaml` is the streaming stage config. +- Streaming is currently single-request oriented; the fixed smoke matrix now lives in `benchmarks/voxcpm/vllm_omni/run_offline_matrix.py`. +- `ref_text` must be the real transcript of the reference audio. Mismatched text usually causes obvious quality degradation. diff --git a/examples/offline_inference/voxcpm/end2end.py b/examples/offline_inference/voxcpm/end2end.py new file mode 100644 index 000000000..980410fea --- /dev/null +++ b/examples/offline_inference/voxcpm/end2end.py @@ -0,0 +1,206 @@ +"""Minimal offline VoxCPM example for vLLM Omni.""" + +from __future__ import annotations + +import asyncio +import time +from pathlib import Path +from typing import Any + +import soundfile as sf +import torch +from vllm.utils.argparse_utils import FlexibleArgumentParser + +from vllm_omni import AsyncOmni, Omni + +REPO_ROOT = Path(__file__).resolve().parents[3] +DEFAULT_SYNC_STAGE_CONFIG = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml" + + +def _build_prompt(args) -> dict[str, Any]: + additional_information: dict[str, list[Any]] = { + "text": [args.text], + "cfg_value": [args.cfg_value], + "inference_timesteps": [args.inference_timesteps], + "min_len": [args.min_len], + "max_new_tokens": [args.max_new_tokens], + } + if args.streaming_prefix_len is not None: + additional_information["streaming_prefix_len"] = [args.streaming_prefix_len] + if args.ref_audio is not None: + additional_information["ref_audio"] = [args.ref_audio] + if args.ref_text is not None: + additional_information["ref_text"] = [args.ref_text] + return { + "prompt_token_ids": [1], + "additional_information": additional_information, + } + + +def _extract_audio_tensor(mm: dict[str, Any]) -> torch.Tensor: + audio = mm.get("audio", mm.get("model_outputs")) + if audio is None: + raise ValueError("No audio output found in multimodal output.") + if isinstance(audio, list): + parts = [torch.as_tensor(item).float().cpu().reshape(-1) for item in audio] + audio = torch.cat(parts, dim=-1) if parts else torch.zeros(0) + if not isinstance(audio, torch.Tensor): + audio = torch.as_tensor(audio) + return audio.float().cpu().reshape(-1) + + +def _extract_sample_rate(mm: dict[str, Any]) -> int: + sr_raw = mm.get("sr", 24000) + if isinstance(sr_raw, list) and sr_raw: + sr_raw = sr_raw[-1] + if hasattr(sr_raw, "item"): + return int(sr_raw.item()) + return int(sr_raw) + + +def _is_streaming_stage_config(stage_config_path: str) -> bool: + return "async_chunk" in Path(stage_config_path).stem + + +def _save_audio(audio: torch.Tensor, sample_rate: int, output_dir: Path, request_id: str) -> Path: + output_dir.mkdir(parents=True, exist_ok=True) + output_path = output_dir / f"output_{request_id}.wav" + sf.write( + output_path, + audio.float().cpu().clamp(-1.0, 1.0).numpy(), + sample_rate, + format="WAV", + subtype="PCM_16", + ) + return output_path + + +async def _run_streaming(args) -> Path: + prompt = _build_prompt(args) + output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio_streaming") + request_id = "streaming_example" + sample_rate = 24000 + buffered_samples = 0 + chunks: list[torch.Tensor] = [] + started = time.perf_counter() + omni = AsyncOmni( + model=args.model, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + try: + async for stage_output in omni.generate(prompt, request_id=request_id): + mm = getattr(stage_output, "multimodal_output", None) + if not isinstance(mm, dict): + request_output = getattr(stage_output, "request_output", None) + if request_output is None: + continue + mm = getattr(request_output, "multimodal_output", None) + if not isinstance(mm, dict) and getattr(request_output, "outputs", None): + mm = getattr(request_output.outputs[0], "multimodal_output", None) + if not isinstance(mm, dict): + continue + audio = _extract_audio_tensor(mm) + if audio.numel() == 0: + continue + sample_rate = _extract_sample_rate(mm) + if audio.numel() > buffered_samples: + delta = audio[buffered_samples:] + buffered_samples = int(audio.numel()) + else: + delta = audio + buffered_samples += int(delta.numel()) + if delta.numel() > 0: + chunks.append(delta) + if not chunks: + raise RuntimeError("No streaming audio chunks received from VoxCPM.") + output_audio = torch.cat(chunks, dim=0) + output_path = _save_audio(output_audio, sample_rate, output_dir, request_id) + print(f"Saved streaming audio to: {output_path} ({time.perf_counter() - started:.2f}s)") + return output_path + finally: + omni.shutdown() + + +def _run_sync(args) -> Path: + prompt = _build_prompt(args) + output_dir = Path(args.output_dir) if args.output_dir is not None else Path("output_audio") + request_id = "sync_example" + started = time.perf_counter() + last_mm: dict[str, Any] | None = None + omni = Omni( + model=args.model, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + stage_init_timeout=args.stage_init_timeout, + ) + for stage_outputs in omni.generate(prompt): + request_output = getattr(stage_outputs, "request_output", None) + if request_output is None: + continue + outputs = getattr(request_output, "outputs", None) + if outputs: + for output in outputs: + mm = getattr(output, "multimodal_output", None) + if isinstance(mm, dict): + last_mm = mm + mm = getattr(request_output, "multimodal_output", None) + if isinstance(mm, dict): + last_mm = mm + if last_mm is None: + raise RuntimeError("No audio output received from VoxCPM.") + output_path = _save_audio( + _extract_audio_tensor(last_mm), + _extract_sample_rate(last_mm), + output_dir, + request_id, + ) + print(f"Saved audio to: {output_path} ({time.perf_counter() - started:.2f}s)") + return output_path + + +def parse_args(): + parser = FlexibleArgumentParser(description="Minimal offline VoxCPM example for vLLM Omni.") + parser.add_argument("--model", type=str, required=True, help="Local VoxCPM model directory.") + parser.add_argument( + "--stage-configs-path", + type=str, + default=str(DEFAULT_SYNC_STAGE_CONFIG), + help=("Stage config path. Use voxcpm.yaml for non-streaming or voxcpm_async_chunk.yaml for streaming."), + ) + parser.add_argument("--text", type=str, required=True, help="Input text for synthesis.") + parser.add_argument("--ref-audio", type=str, default=None, help="Reference audio path for voice cloning.") + parser.add_argument("--ref-text", type=str, default=None, help="Transcript of the reference audio.") + parser.add_argument("--output-dir", type=str, default=None, help="Output directory for generated wav files.") + parser.add_argument("--cfg-value", type=float, default=2.0, help="Guidance value passed to VoxCPM.") + parser.add_argument("--inference-timesteps", type=int, default=10, help="Number of diffusion timesteps.") + parser.add_argument("--min-len", type=int, default=2, help="Minimum latent length.") + parser.add_argument("--max-new-tokens", type=int, default=4096, help="Maximum latent length.") + parser.add_argument( + "--streaming-prefix-len", + type=int, + default=3, + help="Streaming prefix length used by voxcpm_async_chunk.yaml.", + ) + parser.add_argument("--stage-init-timeout", type=int, default=600, help="Stage initialization timeout in seconds.") + parser.add_argument("--log-stats", action="store_true", help="Enable vLLM Omni stats logging.") + args = parser.parse_args() + if (args.ref_audio is None) != (args.ref_text is None): + raise ValueError("Voice cloning requires --ref-audio and --ref-text together.") + return args + + +def main(args) -> None: + route = "streaming" if _is_streaming_stage_config(args.stage_configs_path) else "sync" + print(f"Model: {args.model}") + print(f"Stage config: {args.stage_configs_path}") + print(f"Route: {route}") + if route == "streaming": + asyncio.run(_run_streaming(args)) + else: + _run_sync(args) + + +if __name__ == "__main__": + main(parse_args()) diff --git a/examples/online_serving/voxcpm/README.md b/examples/online_serving/voxcpm/README.md new file mode 100644 index 000000000..78e1bf4aa --- /dev/null +++ b/examples/online_serving/voxcpm/README.md @@ -0,0 +1,166 @@ +# VoxCPM + +## Prerequisites + +Install VoxCPM in one of these ways: + +```bash +pip install voxcpm +``` + +or point vLLM-Omni to a local VoxCPM source tree: + +```bash +export VLLM_OMNI_VOXCPM_CODE_PATH=/path/to/VoxCPM/src +``` + +If the native VoxCPM `config.json` lacks HF metadata such as `model_type`, +prepare a persistent HF-compatible config directory and export: + +```bash +export VLLM_OMNI_VOXCPM_HF_CONFIG_PATH=/tmp/voxcpm_hf_config +mkdir -p "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH" +cp "$VOXCPM_MODEL/config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/config.json" +cp "$VOXCPM_MODEL/generation_config.json" "$VLLM_OMNI_VOXCPM_HF_CONFIG_PATH/generation_config.json" 2>/dev/null || true +python3 -c 'import json, os; p=os.path.join(os.environ["VLLM_OMNI_VOXCPM_HF_CONFIG_PATH"], "config.json"); cfg=json.load(open(p, "r", encoding="utf-8")); cfg["model_type"]="voxcpm"; cfg.setdefault("architectures", ["VoxCPMForConditionalGeneration"]); json.dump(cfg, open(p, "w", encoding="utf-8"), indent=2, ensure_ascii=False)' +``` + +The VoxCPM stage configs read `VLLM_OMNI_VOXCPM_HF_CONFIG_PATH` directly. The `python3 -c` form above avoids heredoc/indentation issues in interactive shells. + +## Launch the Server + +Use the async-chunk stage config by default: + +```bash +export VOXCPM_MODEL=/path/to/voxcpm-model +cd examples/online_serving/voxcpm +./run_server.sh +``` + +Use the non-streaming stage config: + +```bash +./run_server.sh sync +``` + +You can also launch the server directly: + +```bash +vllm serve "$VOXCPM_MODEL" \ + --stage-configs-path vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml \ + --trust-remote-code \ + --enforce-eager \ + --omni \ + --port 8091 +``` + +## Send Requests + +### Basic text-to-speech + +```bash +python openai_speech_client.py \ + --model "$VOXCPM_MODEL" \ + --text "This is a VoxCPM online text-to-speech example." +``` + +### Voice cloning + +```bash +python openai_speech_client.py \ + --model "$VOXCPM_MODEL" \ + --text "This sentence is synthesized with a cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript spoken in reference.wav." +``` + +`ref_text` must be the real transcript of the reference audio. Placeholder text or mismatched text will usually degrade quality badly. + +### Streaming PCM output + +```bash +python openai_speech_client.py \ + --model "$VOXCPM_MODEL" \ + --text "This is a streaming VoxCPM request." \ + --stream \ + --output voxcpm_stream.pcm +``` + +### Using curl + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "OpenBMB/VoxCPM1.5", + "input": "Hello from VoxCPM online serving.", + "response_format": "wav" + }' --output output.wav +``` + +Voice cloning: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "OpenBMB/VoxCPM1.5", + "input": "This sentence uses a cloned voice.", + "ref_audio": "https://example.com/reference.wav", + "ref_text": "The exact transcript spoken in the reference audio.", + "response_format": "wav" + }' --output cloned.wav +``` + +Streaming PCM: + +```bash +curl -X POST http://localhost:8091/v1/audio/speech \ + -H "Content-Type: application/json" \ + -d '{ + "model": "OpenBMB/VoxCPM1.5", + "input": "This is a streaming VoxCPM request.", + "stream": true, + "response_format": "pcm" + }' --output output.pcm +``` + +## Supported Request Shape + +VoxCPM online serving currently supports: + +- plain text-to-speech +- voice cloning with `ref_audio` + `ref_text` +- `stream=true` with `response_format=pcm` or `wav` + +VoxCPM online serving does not use these generic TTS fields: + +- `voice` +- `instructions` +- `language` +- `speaker_embedding` +- `x_vector_only_mode` + +## Streaming vs Non-Streaming + +- `voxcpm_async_chunk.yaml` enables async-chunk streaming and is best for single-request streaming latency. +- `voxcpm.yaml` performs one-shot latent generation then VAE decode. + +Like native VoxCPM, the async streaming path should be treated as single-request. If you need stable throughput benchmarking, prefer `voxcpm.yaml`. + +Do not use `voxcpm_async_chunk.yaml` for concurrent online streaming or `/v1/audio/speech/batch`. For multiple requests, prefer `voxcpm.yaml`. + +## Benchmark + +The serving benchmark reports TTFP and RTF: + +```bash +python benchmarks/voxcpm/vllm_omni/bench_tts_serve.py \ + --host 127.0.0.1 \ + --port 8091 \ + --num-prompts 10 \ + --max-concurrency 1 \ + --result-dir /tmp/voxcpm_bench +``` + +For the async-chunk server, keep `--max-concurrency 1`. diff --git a/examples/online_serving/voxcpm/openai_speech_client.py b/examples/online_serving/voxcpm/openai_speech_client.py new file mode 100644 index 000000000..c400114e8 --- /dev/null +++ b/examples/online_serving/voxcpm/openai_speech_client.py @@ -0,0 +1,155 @@ +"""OpenAI-compatible client for VoxCPM via /v1/audio/speech. + +Examples: + # Basic text-to-speech + python openai_speech_client.py --text "Hello from VoxCPM" + + # Voice cloning + python openai_speech_client.py \ + --text "This sentence uses the cloned voice." \ + --ref-audio /path/to/reference.wav \ + --ref-text "The exact transcript spoken in the reference audio." + + # Streaming PCM output + python openai_speech_client.py \ + --text "This is a streaming VoxCPM request." \ + --stream \ + --output output.pcm +""" + +import argparse +import base64 +import os + +import httpx + +DEFAULT_API_BASE = "http://localhost:8091" +DEFAULT_API_KEY = "EMPTY" +DEFAULT_MODEL = "OpenBMB/VoxCPM1.5" + + +def encode_audio_to_base64(audio_path: str) -> str: + """Encode a local audio file to base64 data URL.""" + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + + ext = audio_path.lower().rsplit(".", 1)[-1] + mime_map = { + "wav": "audio/wav", + "mp3": "audio/mpeg", + "flac": "audio/flac", + "ogg": "audio/ogg", + } + mime_type = mime_map.get(ext, "audio/wav") + + with open(audio_path, "rb") as f: + audio_b64 = base64.b64encode(f.read()).decode("utf-8") + return f"data:{mime_type};base64,{audio_b64}" + + +def build_payload(args) -> dict[str, object]: + payload: dict[str, object] = { + "model": args.model, + "input": args.text, + "response_format": "pcm" if args.stream else args.response_format, + } + + if args.ref_audio: + if args.ref_audio.startswith(("http://", "https://", "data:")): + payload["ref_audio"] = args.ref_audio + else: + payload["ref_audio"] = encode_audio_to_base64(args.ref_audio) + if args.ref_text: + payload["ref_text"] = args.ref_text + if args.max_new_tokens is not None: + payload["max_new_tokens"] = args.max_new_tokens + if args.stream: + payload["stream"] = True + + return payload + + +def run_tts(args) -> None: + payload = build_payload(args) + api_url = f"{args.api_base}/v1/audio/speech" + headers = { + "Content-Type": "application/json", + "Authorization": f"Bearer {args.api_key}", + } + + print(f"Model: {args.model}") + print(f"Text: {args.text}") + if args.ref_audio: + print("Mode: voice cloning") + print(f"Reference audio: {args.ref_audio}") + else: + print("Mode: text-to-speech") + + if args.stream: + output_path = args.output or "voxcpm_output.pcm" + with httpx.Client(timeout=300.0) as client: + with client.stream("POST", api_url, json=payload, headers=headers) as response: + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.read().decode("utf-8", errors="ignore")) + return + + total_bytes = 0 + with open(output_path, "wb") as f: + for chunk in response.iter_bytes(): + if not chunk: + continue + f.write(chunk) + total_bytes += len(chunk) + print(f"Streamed {total_bytes} bytes to: {output_path}") + return + + with httpx.Client(timeout=300.0) as client: + response = client.post(api_url, json=payload, headers=headers) + + if response.status_code != 200: + print(f"Error: {response.status_code}") + print(response.text) + return + + try: + text = response.content.decode("utf-8") + if text.startswith('{"error"'): + print(f"Error: {text}") + return + except UnicodeDecodeError: + pass + + output_path = args.output or "voxcpm_output.wav" + with open(output_path, "wb") as f: + f.write(response.content) + print(f"Audio saved to: {output_path}") + + +def main(): + parser = argparse.ArgumentParser(description="VoxCPM OpenAI-compatible speech client") + parser.add_argument("--api-base", default=DEFAULT_API_BASE, help="API base URL") + parser.add_argument("--api-key", default=DEFAULT_API_KEY, help="API key") + parser.add_argument("--model", "-m", default=DEFAULT_MODEL, help="Model name or path") + parser.add_argument("--text", required=True, help="Text to synthesize") + parser.add_argument("--ref-audio", default=None, help="Reference audio path, URL, or data URL") + parser.add_argument( + "--ref-text", + default=None, + help="The exact transcript spoken in the reference audio", + ) + parser.add_argument("--stream", action="store_true", help="Enable streaming PCM output") + parser.add_argument( + "--response-format", + default="wav", + choices=["wav", "pcm", "flac", "mp3", "aac", "opus"], + help="Audio format for non-streaming mode (default: wav)", + ) + parser.add_argument("--max-new-tokens", type=int, default=None, help="Maximum tokens to generate") + parser.add_argument("--output", "-o", default=None, help="Output file path") + args = parser.parse_args() + run_tts(args) + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/voxcpm/run_server.sh b/examples/online_serving/voxcpm/run_server.sh new file mode 100755 index 000000000..ab4b6fe85 --- /dev/null +++ b/examples/online_serving/voxcpm/run_server.sh @@ -0,0 +1,38 @@ +#!/bin/bash +# Launch vLLM-Omni server for VoxCPM online speech serving. +# +# Usage: +# ./run_server.sh # default: async_chunk stage config +# ./run_server.sh async # async_chunk stage config +# ./run_server.sh sync # no-async-chunk stage config +# VOXCPM_MODEL=/path/to/model ./run_server.sh + +set -e + +MODE="${1:-async}" +MODEL="${VOXCPM_MODEL:-OpenBMB/VoxCPM1.5}" + +case "$MODE" in + async) + STAGE_CONFIG="vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml" + ;; + sync) + STAGE_CONFIG="vllm_omni/model_executor/stage_configs/voxcpm.yaml" + ;; + *) + echo "Unknown mode: $MODE" + echo "Supported: async, sync" + exit 1 + ;; +esac + +echo "Starting VoxCPM server with model: $MODEL" +echo "Stage config: $STAGE_CONFIG" + +vllm serve "$MODEL" \ + --stage-configs-path "$STAGE_CONFIG" \ + --host 0.0.0.0 \ + --port 8091 \ + --trust-remote-code \ + --enforce-eager \ + --omni diff --git a/tests/e2e/offline_inference/test_voxcpm.py b/tests/e2e/offline_inference/test_voxcpm.py new file mode 100644 index 000000000..d7f65525e --- /dev/null +++ b/tests/e2e/offline_inference/test_voxcpm.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""E2E test for VoxCPM offline inference.""" + +import json +import os +from pathlib import Path +from typing import Any + +import numpy as np +import pytest +import torch + +import tests.conftest as omni_test_conftest +from tests.conftest import OmniRunner +from tests.utils import hardware_test +from vllm_omni.model_executor.models.voxcpm.voxcpm_runtime_utils import ( + prepare_voxcpm_hf_config_dir, + resolve_voxcpm_model_dir, +) + +VOXCPM_MODEL = os.environ.get("VOXCPM_MODEL", "OpenBMB/VoxCPM1.5") +STAGE_CONFIG = str( + Path(__file__).parent.parent.parent.parent / "vllm_omni" / "model_executor" / "stage_configs" / "voxcpm.yaml" +) +SAMPLE_RATE = 24000 + + +@pytest.fixture(autouse=True) +def _patch_npu_cleanup_for_voxcpm(monkeypatch: pytest.MonkeyPatch): + """Limit the NPU cleanup workaround to this VoxCPM test module only.""" + original_cleanup = omni_test_conftest.cleanup_dist_env_and_memory + + def _safe_cleanup() -> None: + try: + original_cleanup() + except RuntimeError as exc: + if "Allocator for npu is not a DeviceAllocator" in str(exc): + return + raise + + monkeypatch.setattr(omni_test_conftest, "cleanup_dist_env_and_memory", _safe_cleanup) + + +def _build_prompt(text: str) -> dict[str, Any]: + return { + "prompt_token_ids": [1], + "additional_information": { + "text": [text], + "cfg_value": [2.0], + "inference_timesteps": [10], + "min_len": [2], + "max_new_tokens": [1024], + }, + } + + +def _extract_audio_tensor(multimodal_output: dict[str, Any]) -> torch.Tensor: + audio = multimodal_output.get("audio", multimodal_output.get("model_outputs")) + assert audio is not None, f"No audio output found, keys={list(multimodal_output.keys())}" + + if isinstance(audio, list): + parts: list[torch.Tensor] = [] + for item in audio: + if item is None: + continue + tensor = torch.as_tensor(item) + if tensor.numel() == 0: + continue + parts.append(tensor.float().cpu().reshape(-1)) + return torch.cat(parts, dim=-1) if parts else torch.zeros((0,), dtype=torch.float32) + + return torch.as_tensor(audio).float().cpu().reshape(-1) + + +def _extract_final_multimodal_output(outputs) -> dict[str, Any]: + for item in reversed(outputs): + request_output = getattr(item, "request_output", None) + if request_output is not None: + multimodal_output = getattr(request_output, "multimodal_output", None) + if isinstance(multimodal_output, dict): + return multimodal_output + completions = getattr(request_output, "outputs", None) or [] + for completion in completions: + multimodal_output = getattr(completion, "multimodal_output", None) + if isinstance(multimodal_output, dict): + return multimodal_output + + multimodal_output = getattr(item, "multimodal_output", None) + if isinstance(multimodal_output, dict): + return multimodal_output + + raise AssertionError("No multimodal audio output found in VoxCPM generate results") + + +@pytest.fixture +def voxcpm_model_path(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> str: + model_dir = resolve_voxcpm_model_dir(VOXCPM_MODEL) + + hf_config_env = os.environ.get("VLLM_OMNI_VOXCPM_HF_CONFIG_PATH") + if hf_config_env: + hf_config_dir = Path(hf_config_env).expanduser() + else: + hf_config_dir = tmp_path / "voxcpm_hf_config" + + if not (hf_config_dir / "config.json").exists(): + prepare_voxcpm_hf_config_dir(model_dir, hf_config_dir) + + monkeypatch.setenv("VLLM_OMNI_VOXCPM_HF_CONFIG_PATH", str(hf_config_dir)) + return str(model_dir) + + +def test_prepare_voxcpm_hf_config_dir(tmp_path: Path): + model_dir = tmp_path / "model" + model_dir.mkdir() + (model_dir / "config.json").write_text(json.dumps({"hidden_size": 1024}), encoding="utf-8") + (model_dir / "generation_config.json").write_text(json.dumps({"do_sample": False}), encoding="utf-8") + + hf_config_dir = prepare_voxcpm_hf_config_dir(model_dir, tmp_path / "voxcpm_hf_config") + + prepared_config = json.loads((hf_config_dir / "config.json").read_text(encoding="utf-8")) + assert prepared_config["model_type"] == "voxcpm" + assert prepared_config["architectures"] == ["VoxCPMForConditionalGeneration"] + assert (hf_config_dir / "generation_config.json").exists() + + +def test_resolve_voxcpm_model_dir_local_path(tmp_path: Path): + model_dir = tmp_path / "OpenBMB" / "VoxCPM1.5" + model_dir.mkdir(parents=True) + + assert resolve_voxcpm_model_dir(str(model_dir)) == model_dir + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "L4"}, num_cards=1) +def test_voxcpm_zero_shot_001(voxcpm_model_path: str): + with OmniRunner(voxcpm_model_path, stage_configs_path=STAGE_CONFIG) as runner: + outputs = list(runner.omni.generate(_build_prompt("Hello, this is a VoxCPM offline inference test."))) + + assert outputs, "No outputs returned" + + multimodal_output = _extract_final_multimodal_output(outputs) + audio = _extract_audio_tensor(multimodal_output) + assert audio.numel() > SAMPLE_RATE // 2, f"Audio too short: {audio.numel()} samples" + + duration_s = audio.shape[0] / SAMPLE_RATE + assert 0.5 < duration_s < 30.0, f"Audio duration out of range: {duration_s:.2f}s" + + peak = float(torch.max(torch.abs(audio)).item()) if audio.numel() > 0 else 0.0 + assert peak > 0.01, "Generated audio appears to be silence" + + audio_np = audio.numpy() + rms = float(np.sqrt(np.mean(np.square(audio_np)))) if audio_np.size else 0.0 + assert rms > 1e-4, "Generated audio RMS too low" diff --git a/tests/engine/test_arg_utils.py b/tests/engine/test_arg_utils.py index 35d55f1cc..565c83c1a 100644 --- a/tests/engine/test_arg_utils.py +++ b/tests/engine/test_arg_utils.py @@ -7,6 +7,7 @@ explicitly patch values that differ from vLLM. import argparse import inspect from types import SimpleNamespace +from unittest.mock import Mock import pytest from pydantic import ValidationError @@ -166,6 +167,24 @@ def test_stage_configs_path_field(): assert args.stage_configs_path == "/some/path.yaml" +def test_voxcpm_model_arch_injects_model_type_override(mocker): + """Ensure VoxCPM model_arch injects hf_overrides for config resolution.""" + mocker.patch.object(OmniEngineArgs, "_ensure_omni_models_registered", return_value=True) + mocker.patch.object(OmniEngineArgs, "_patch_empty_hf_config") + mocker.patch.object(EngineArgs, "create_model_config", return_value=Mock()) + mocker.patch.object(OmniModelConfig, "from_vllm_model_config", return_value=Mock()) + + args = OmniEngineArgs( + model="OpenBMB/VoxCPM1.5", + model_arch="VoxCPMForConditionalGeneration", + ) + args.create_model_config() + + assert args.hf_overrides["architectures"] == ["VoxCPMForConditionalGeneration"] + assert args.hf_overrides["model_type"] == "voxcpm" + args._patch_empty_hf_config.assert_called_once_with("voxcpm") + + def test_strip_single_engine_args(): """_strip_single_engine_args should remove EngineArgs fields but keep omni fields.""" kwargs = { diff --git a/tests/entrypoints/openai_api/test_serving_speech_voxcpm.py b/tests/entrypoints/openai_api/test_serving_speech_voxcpm.py new file mode 100644 index 000000000..48660b6d1 --- /dev/null +++ b/tests/entrypoints/openai_api/test_serving_speech_voxcpm.py @@ -0,0 +1,143 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""UTs for VoxCPM OpenAI speech serving behavior.""" + +import asyncio +from types import SimpleNamespace +from unittest.mock import AsyncMock + +import pytest +from pytest_mock import MockerFixture + +from vllm_omni.entrypoints.openai.protocol.audio import OpenAICreateSpeechRequest +from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +@pytest.fixture +def voxcpm_server(mocker: MockerFixture): + mocker.patch.object(OmniOpenAIServingSpeech, "_load_supported_speakers", return_value=set()) + mocker.patch.object(OmniOpenAIServingSpeech, "_load_codec_frame_rate", return_value=None) + + mock_engine_client = mocker.MagicMock() + mock_engine_client.errored = False + mock_engine_client.model_config = mocker.MagicMock(model="OpenBMB/VoxCPM1.5") + mock_engine_client.default_sampling_params_list = [SimpleNamespace(max_tokens=2048)] + mock_engine_client.tts_batch_max_items = 32 + mock_engine_client.generate = mocker.MagicMock(return_value="generator") + mock_engine_client.stage_configs = [ + SimpleNamespace( + engine_args=SimpleNamespace( + model_stage="latent_generator", + model_arch="VoxCPMForConditionalGeneration", + ), + tts_args={}, + ), + SimpleNamespace( + engine_args=SimpleNamespace(model_stage="vae"), + tts_args={}, + ), + ] + + mock_models = mocker.MagicMock() + mock_models.is_base_model.return_value = True + + return OmniOpenAIServingSpeech( + engine_client=mock_engine_client, + models=mock_models, + request_logger=mocker.MagicMock(), + ) + + +class TestVoxCPMServing: + def test_voxcpm_model_type_detection(self, voxcpm_server): + assert voxcpm_server._tts_model_type == "voxcpm" + assert voxcpm_server._is_tts is True + assert voxcpm_server.supported_speakers == set() + + @pytest.mark.parametrize( + ("request_kwargs", "expected_substring"), + [ + ({"voice": "alice"}, "voice"), + ({"instructions": "whisper"}, "instructions"), + ({"language": "en"}, "language"), + ({"task_type": "CustomVoice"}, "plain tts"), + ({"x_vector_only_mode": True}, "x_vector_only_mode"), + ({"speaker_embedding": [0.1, 0.2]}, "speaker_embedding"), + ({"initial_codec_chunk_frames": 4}, "initial_codec_chunk_frames"), + ({"ref_text": "reference"}, "ref_audio"), + ], + ) + def test_validate_voxcpm_rejects_unsupported_fields(self, voxcpm_server, request_kwargs, expected_substring): + request = OpenAICreateSpeechRequest(input="hello voxcpm", **request_kwargs) + error = voxcpm_server._validate_voxcpm_request(request) + assert error is not None + assert expected_substring in error.lower() + + def test_validate_voxcpm_accepts_plain_tts_request(self, voxcpm_server): + request = OpenAICreateSpeechRequest(input="hello voxcpm", max_new_tokens=256) + assert voxcpm_server._validate_voxcpm_request(request) is None + + def test_validate_voxcpm_accepts_voice_clone_request(self, voxcpm_server): + request = OpenAICreateSpeechRequest( + input="clone this voice", + ref_audio="data:audio/wav;base64,QUJD", + ref_text="reference transcript", + max_new_tokens=256, + ) + assert voxcpm_server._validate_voxcpm_request(request) is None + + def test_prepare_speech_generation_voxcpm_text_only(self, voxcpm_server): + request = OpenAICreateSpeechRequest(input="hello voxcpm", max_new_tokens=321) + + request_id, generator, tts_params = asyncio.run(voxcpm_server._prepare_speech_generation(request)) + + assert request_id.startswith("speech-") + assert generator == "generator" + assert tts_params == { + "text": ["hello voxcpm"], + "cfg_value": [2.0], + "inference_timesteps": [10], + "min_len": [2], + "max_new_tokens": [321], + } + + voxcpm_server.engine_client.generate.assert_called_once() + call = voxcpm_server.engine_client.generate.call_args + assert call.kwargs["prompt"] == { + "prompt_token_ids": [1], + "additional_information": tts_params, + } + assert call.kwargs["output_modalities"] == ["audio"] + + def test_prepare_speech_generation_voxcpm_voice_clone_resolves_ref_audio(self, voxcpm_server): + voxcpm_server._resolve_ref_audio = AsyncMock(return_value=([0.1, -0.1, 0.2], 16000)) + request = OpenAICreateSpeechRequest( + input="clone this voice", + ref_audio="data:audio/wav;base64,QUJD", + ref_text="reference transcript", + max_new_tokens=512, + ) + + request_id, generator, tts_params = asyncio.run(voxcpm_server._prepare_speech_generation(request)) + + assert request_id.startswith("speech-") + assert generator == "generator" + assert tts_params == { + "text": ["clone this voice"], + "cfg_value": [2.0], + "inference_timesteps": [10], + "min_len": [2], + "max_new_tokens": [512], + "ref_text": ["reference transcript"], + "ref_audio": [[[0.1, -0.1, 0.2], 16000]], + } + + voxcpm_server._resolve_ref_audio.assert_awaited_once_with("data:audio/wav;base64,QUJD") + call = voxcpm_server.engine_client.generate.call_args + assert call.kwargs["prompt"] == { + "prompt_token_ids": [1], + "additional_information": tts_params, + } diff --git a/tests/entrypoints/test_utils.py b/tests/entrypoints/test_utils.py index 94e254c25..248629d51 100644 --- a/tests/entrypoints/test_utils.py +++ b/tests/entrypoints/test_utils.py @@ -310,6 +310,39 @@ class TestResolveModelConfigPath: assert result is not None assert "glm_image.yaml" in result + def test_voxcpm_transformers_format_resolution(self, mocker: MockerFixture): + """Test VoxCPM transformers config resolves to the voxcpm stage config.""" + mocker.patch( + "vllm_omni.entrypoints.utils.get_config", + side_effect=ValueError("missing transformers config"), + ) + mocker.patch( + "vllm_omni.entrypoints.utils.file_or_path_exists", + side_effect=lambda _model, filename, revision=None: filename == "config.json", + ) + mocker.patch( + "vllm_omni.entrypoints.utils.get_hf_file_to_dict", + return_value={"model_type": "voxcpm"}, + ) + mocker.patch( + "vllm_omni.entrypoints.utils.current_omni_platform.get_default_stage_config_path", + return_value="vllm_omni/model_executor/stage_configs", + ) + + original_exists = os.path.exists + + def mock_exists(path): + if "voxcpm.yaml" in str(path): + return True + return original_exists(path) + + mocker.patch("os.path.exists", side_effect=mock_exists) + + result = resolve_model_config_path("OpenBMB/VoxCPM1.5") + + assert result is not None + assert "voxcpm.yaml" in result + class TestLoadAndResolveStageConfigs: def test_load_and_resolve_with_kwargs(self): diff --git a/tests/model_executor/stage_input_processors/test_voxcpm_async_chunk.py b/tests/model_executor/stage_input_processors/test_voxcpm_async_chunk.py new file mode 100644 index 000000000..7d6fc6e74 --- /dev/null +++ b/tests/model_executor/stage_input_processors/test_voxcpm_async_chunk.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""UTs for VoxCPM async-chunk stage input processing.""" + +from types import SimpleNamespace + +import pytest +import torch + +from vllm_omni.model_executor.stage_input_processors.voxcpm import ( + _VOXCPM_LATENT_MAGIC, + _coerce_finished_flag, + latent2vae_async_chunk, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _request(*, finished): + return SimpleNamespace(is_finished=lambda: finished) + + +def _decode_serialized_latent(codes: list[int]) -> torch.Tensor: + assert codes[0] == _VOXCPM_LATENT_MAGIC + latent_dim = codes[1] + time_dim = codes[2] + payload = torch.tensor(codes[3:], dtype=torch.int32).to(torch.uint16) + return payload.view(torch.bfloat16).to(torch.float32).reshape(1, latent_dim, time_dim) + + +@pytest.mark.parametrize( + ("value", "expected"), + [ + (None, False), + (False, False), + (True, True), + (torch.tensor(False), False), + (torch.tensor(True), True), + ([torch.tensor(True)], True), + (([True],), True), + ([], False), + ], +) +def test_coerce_finished_flag(value, expected): + assert _coerce_finished_flag(value) is expected + + +def test_latent2vae_async_chunk_serializes_latent_payload(): + latent = torch.arange(6, dtype=torch.float32).reshape(2, 3) + + payload = latent2vae_async_chunk( + transfer_manager=None, + pooling_output={"latent_audio_feat": latent}, + request=_request(finished=False), + is_finished=torch.tensor(False), + ) + + assert payload is not None + assert torch.equal(payload["finished"], torch.tensor(False, dtype=torch.bool)) + recovered = _decode_serialized_latent(payload["code_predictor_codes"]) + torch.testing.assert_close(recovered, latent.to(torch.bfloat16).to(torch.float32).unsqueeze(0)) + + +def test_latent2vae_async_chunk_returns_terminal_marker_without_latent(): + payload = latent2vae_async_chunk( + transfer_manager=None, + pooling_output=None, + request=_request(finished=[torch.tensor(True)]), + is_finished=False, + ) + + assert payload == { + "code_predictor_codes": [], + "finished": torch.tensor(True, dtype=torch.bool), + } + + +def test_latent2vae_async_chunk_returns_none_for_nonterminal_empty_chunk(): + payload = latent2vae_async_chunk( + transfer_manager=None, + pooling_output={"latent_audio_feat": torch.zeros((0,), dtype=torch.float32)}, + request=_request(finished=False), + is_finished=False, + ) + + assert payload is None diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index d61102c7e..5b69d6b1f 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -21,6 +21,7 @@ _ARCH_TO_MODEL_TYPE: dict[str, str] = { "CosyVoice3Model": "cosyvoice3", "OmniVoiceModel": "omnivoice", "VoxCPM2TalkerForConditionalGeneration": "voxcpm2", + "VoxCPMForConditionalGeneration": "voxcpm", } # Maps model architecture names to tokenizer subfolder paths within HF repos. @@ -41,6 +42,7 @@ def _register_omni_hf_configs() -> None: from vllm_omni.model_executor.models.voxtral_tts.configuration_voxtral_tts import ( VoxtralTTSConfig, ) + from vllm_omni.transformers_utils.configs.voxcpm import VoxCPMConfig from vllm_omni.transformers_utils.configs.voxcpm2 import VoxCPM2Config except Exception as exc: # pragma: no cover - best-effort optional registration logger.warning("Skipping omni HF config registration due to import error: %s", exc) @@ -59,6 +61,7 @@ def _register_omni_hf_configs() -> None: ("cosyvoice3", CosyVoice3Config), ("omnivoice", OmniVoiceConfig), ("voxtral_tts", VoxtralTTSConfig), + ("voxcpm", VoxCPMConfig), ("voxcpm2", VoxCPM2Config), ]: try: diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 1d9754853..1f78f5691 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -49,6 +49,7 @@ _QWEN3_TTS_MODEL_STAGES = {"qwen3_tts"} _FISH_TTS_MODEL_STAGES = {"fish_speech_slow_ar"} _COSYVOICE3_TTS_MODEL_STAGES = {"cosyvoice3_talker"} _OMNIVOICE_TTS_MODEL_STAGES = {"omnivoice_generator"} +_VOXCPM_TTS_MODEL_STAGES = {"latent_generator", "vae"} _VOXCPM2_TTS_MODEL_STAGES = {"latent_generator"} _TTS_MODEL_STAGES: set[str] = ( _VOXTRAL_TTS_MODEL_STAGES @@ -56,6 +57,7 @@ _TTS_MODEL_STAGES: set[str] = ( | _FISH_TTS_MODEL_STAGES | _COSYVOICE3_TTS_MODEL_STAGES | _OMNIVOICE_TTS_MODEL_STAGES + | _VOXCPM_TTS_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES ) _TTS_LANGUAGES: set[str] = { @@ -282,6 +284,11 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if self._tts_stage is None: return None model_stage = getattr(self._tts_stage.engine_args, "model_stage", None) + model_arch = getattr(self._tts_stage.engine_args, "model_arch", None) + if model_arch == "VoxCPM2TalkerForConditionalGeneration": + return "voxcpm2" + if model_arch == "VoxCPMForConditionalGeneration": + return "voxcpm" if model_stage in _QWEN3_TTS_MODEL_STAGES: return "qwen3_tts" if model_stage in _VOXTRAL_TTS_MODEL_STAGES: @@ -292,8 +299,12 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return "cosyvoice3" if model_stage in _OMNIVOICE_TTS_MODEL_STAGES: return "omnivoice" - if model_stage in _VOXCPM2_TTS_MODEL_STAGES: - return "voxcpm2" + if model_stage in (_VOXCPM_TTS_MODEL_STAGES | _VOXCPM2_TTS_MODEL_STAGES): + has_vae_stage = any( + getattr(getattr(stage, "engine_args", None), "model_stage", None) == "vae" + for stage in self.engine_client.stage_configs + ) + return "voxcpm" if has_vae_stage or model_stage == "vae" else "voxcpm2" return None def _compute_max_instructions_length(self) -> int: @@ -318,6 +329,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _load_supported_speakers(self) -> set[str]: """Load supported speakers (case-insensitive) from the model configuration.""" try: + if self._tts_model_type == "voxcpm": + return set() if self._tts_model_type == "voxtral_tts": config = self.engine_client.model_config.hf_config.audio_config else: @@ -377,6 +390,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): def _estimate_prompt_len(self, tts_params: dict[str, Any]) -> int: """Estimate prompt length so the placeholder matches model-side embeddings.""" try: + if self._tts_model_type == "voxcpm": + return 1 from vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_talker import ( Qwen3TTSTalkerForConditionalGeneration, ) @@ -791,6 +806,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return self._validate_fish_tts_request(request) if self._tts_model_type == "cosyvoice3": return self._validate_cosyvoice3_request(request) + if self._tts_model_type == "voxcpm": + return self._validate_voxcpm_request(request) if self._tts_model_type == "voxcpm2": return None # VoxCPM2 accepts any text input return self._validate_qwen_tts_request(request) @@ -832,6 +849,43 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return None + def _validate_voxcpm_request(self, request: OpenAICreateSpeechRequest) -> str | None: + """Validate VoxCPM request parameters. Returns error message or None.""" + if not request.input or not request.input.strip(): + return "Input text cannot be empty" + + if request.voice is not None: + return "'voice' is not supported for VoxCPM" + if request.instructions is not None: + return "'instructions' is not supported for VoxCPM" + if request.language is not None: + return "'language' is not supported for VoxCPM" + if request.task_type not in (None, "Base"): + return "VoxCPM only supports plain TTS or voice cloning with ref_audio/ref_text" + if request.x_vector_only_mode is not None: + return "'x_vector_only_mode' is not supported for VoxCPM" + if request.speaker_embedding is not None: + return "'speaker_embedding' is not supported for VoxCPM" + if request.initial_codec_chunk_frames is not None: + return "'initial_codec_chunk_frames' is not supported for VoxCPM" + + if request.ref_audio is not None: + fmt_err = self._validate_ref_audio_format(request.ref_audio) + if fmt_err: + return fmt_err + if not request.ref_text or not request.ref_text.strip(): + return "Voice cloning requires 'ref_text' (transcript of the reference audio)" + elif request.ref_text is not None: + return "'ref_text' requires 'ref_audio' for VoxCPM voice cloning" + + if request.max_new_tokens is not None: + if request.max_new_tokens < _TTS_MAX_NEW_TOKENS_MIN: + return f"max_new_tokens must be at least {_TTS_MAX_NEW_TOKENS_MIN}" + if request.max_new_tokens > _TTS_MAX_NEW_TOKENS_MAX: + return f"max_new_tokens cannot exceed {_TTS_MAX_NEW_TOKENS_MAX}" + + return None + def _validate_qwen_tts_request(self, request: OpenAICreateSpeechRequest) -> str | None: """Validate Qwen TTS request parameters. Returns error message or None.""" # Infer Base task when ref_audio or ref_text is provided without explicit task_type. @@ -1169,6 +1223,18 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): Processes each parameter if present, skips if not. Values are wrapped in lists as required by the model. """ + if self._tts_model_type == "voxcpm": + params: dict[str, Any] = { + "text": [request.input], + "cfg_value": [2.0], + "inference_timesteps": [10], + "min_len": [2], + "max_new_tokens": [request.max_new_tokens or 4096], + } + if request.ref_text is not None: + params["ref_text"] = [request.ref_text] + return params + params: dict[str, Any] = {} # Text content (always required) @@ -1499,6 +1565,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): model_type = "voxtral_tts" elif self._tts_model_type == "cosyvoice3": model_type = "cosyvoice3" + elif self._tts_model_type == "voxcpm": + model_type = "voxcpm" elif self._tts_model_type == "voxcpm2": model_type = "voxcpm2" elif self._is_tts: diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 089408800..3407b4286 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -145,6 +145,12 @@ _OMNI_MODELS = { "fish_speech_dac_decoder", "FishSpeechDACDecoder", ), + ## VoxCPM + "VoxCPMForConditionalGeneration": ( + "voxcpm", + "voxcpm", + "VoxCPMForConditionalGeneration", + ), ## VoxCPM2 "VoxCPM2TalkerForConditionalGeneration": ( "voxcpm2", diff --git a/vllm_omni/model_executor/models/voxcpm/__init__.py b/vllm_omni/model_executor/models/voxcpm/__init__.py new file mode 100644 index 000000000..3b064c0f6 --- /dev/null +++ b/vllm_omni/model_executor/models/voxcpm/__init__.py @@ -0,0 +1,7 @@ +from .configuration_voxcpm import VoxCPMConfig +from .voxcpm import VoxCPMForConditionalGeneration + +__all__ = [ + "VoxCPMConfig", + "VoxCPMForConditionalGeneration", +] diff --git a/vllm_omni/model_executor/models/voxcpm/configuration_voxcpm.py b/vllm_omni/model_executor/models/voxcpm/configuration_voxcpm.py new file mode 100644 index 000000000..ce1d809bd --- /dev/null +++ b/vllm_omni/model_executor/models/voxcpm/configuration_voxcpm.py @@ -0,0 +1,3 @@ +from vllm_omni.transformers_utils.configs.voxcpm import VoxCPMConfig + +__all__ = ["VoxCPMConfig"] diff --git a/vllm_omni/model_executor/models/voxcpm/voxcpm.py b/vllm_omni/model_executor/models/voxcpm/voxcpm.py new file mode 100644 index 000000000..6fa36fc42 --- /dev/null +++ b/vllm_omni/model_executor/models/voxcpm/voxcpm.py @@ -0,0 +1,886 @@ +from __future__ import annotations + +import json +import os +import sys +import tempfile +import warnings +import wave +from collections.abc import Callable, Generator, Iterable +from pathlib import Path +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +from einops import rearrange +from tqdm import tqdm +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.models.output_templates import OmniOutput + +from .voxcpm_loader import ( + _build_prompt_cache_with_soundfile, + _device_to_string, + _force_cuda_available_for_npu, + _import_voxcpm_audio_vae_classes, + _import_voxcpm_base_model_class, + _is_torchcodec_load_error, + _normalize_dtype_name, + _prepare_runtime_model_dir, + _resolve_runtime_device, +) +from .voxcpm_runtime_utils import resolve_voxcpm_model_dir +from .voxcpm_stage_wrappers import _DirectVoxCPMAudioVAE, _DirectVoxCPMLatentGenerator + +logger = init_logger(__name__) +_VOXCPM_LATENT_MAGIC = 131071 + + +def _make_voxcpm_model_for_omni(base: type[Any]) -> type[Any]: + """Subclass upstream VoxCPMModel: local ``_inference`` + ``latents_only`` prompt-cache generation.""" + + from voxcpm.model.utils import get_dtype + + class VoxCPMModelForOmni(base): + @torch.inference_mode() + def build_prompt_cache(self, *args: Any, **kwargs: Any): + try: + return super().build_prompt_cache(*args, **kwargs) + except (ImportError, ModuleNotFoundError, RuntimeError) as exc: + if not _is_torchcodec_load_error(exc): + raise + return _build_prompt_cache_with_soundfile(self, *args, **kwargs) + + @torch.inference_mode() + def _inference( + self, + text: torch.Tensor, + text_mask: torch.Tensor, + feat: torch.Tensor, + feat_mask: torch.Tensor, + min_len: int = 2, + max_len: int = 2000, + inference_timesteps: int = 10, + cfg_value: float = 2.0, + streaming: bool = False, + streaming_prefix_len: int = 3, + ) -> Generator[tuple[torch.Tensor, torch.Tensor | list[torch.Tensor]], None, None]: + B, _, _, _ = feat.shape + + feat_embed = self.feat_encoder(feat) + feat_embed = self.enc_to_lm_proj(feat_embed) + + scale_emb = self.config.lm_config.scale_emb if self.config.lm_config.use_mup else 1.0 + text_embed = self.base_lm.embed_tokens(text) * scale_emb + combined_embed = text_mask.unsqueeze(-1) * text_embed + feat_mask.unsqueeze(-1) * feat_embed + + prefix_feat_cond = feat[:, -1, ...] + pred_feat_seq: list[torch.Tensor] = [] + + audio_patch_count = int(feat_mask.sum().item()) + if audio_patch_count > 0: + context_len = min(streaming_prefix_len - 1, audio_patch_count) + prompt_context_patches = list(feat[:, -context_len:, :, :].split(1, dim=1)) + pred_feat_seq = prompt_context_patches + pred_feat_seq + + enc_outputs, kv_cache_tuple = self.base_lm( + inputs_embeds=combined_embed, + is_causal=True, + ) + self.base_lm.kv_cache.fill_caches(kv_cache_tuple) + + enc_outputs = self.fsq_layer(enc_outputs) * feat_mask.unsqueeze(-1) + enc_outputs * text_mask.unsqueeze(-1) + lm_hidden = enc_outputs[:, -1, :] + + residual_enc_outputs, residual_kv_cache_tuple = self.residual_lm( + inputs_embeds=enc_outputs + feat_mask.unsqueeze(-1) * feat_embed, + is_causal=True, + ) + self.residual_lm.kv_cache.fill_caches(residual_kv_cache_tuple) + residual_hidden = residual_enc_outputs[:, -1, :] + + for step_idx in tqdm(range(max_len)): + dit_hidden = self.lm_to_dit_proj(lm_hidden) + self.res_to_dit_proj(residual_hidden) + pred_feat = self.feat_decoder( + mu=dit_hidden, + patch_size=self.patch_size, + cond=prefix_feat_cond.transpose(1, 2).contiguous(), + n_timesteps=inference_timesteps, + cfg_value=cfg_value, + ).transpose(1, 2) + + curr_embed = self.enc_to_lm_proj(self.feat_encoder(pred_feat.unsqueeze(1))) + pred_feat_seq.append(pred_feat.unsqueeze(1)) + prefix_feat_cond = pred_feat + + if streaming: + pred_feat_chunk = torch.cat(pred_feat_seq[-streaming_prefix_len:], dim=1) + feat_pred = rearrange(pred_feat_chunk, "b t p d -> b d (t p)", b=B, p=self.patch_size) + yield feat_pred, pred_feat_seq + + stop_flag = self.stop_head(self.stop_actn(self.stop_proj(lm_hidden))).argmax(dim=-1)[0].cpu().item() + if step_idx > min_len and stop_flag == 1: + break + + lm_hidden = self.base_lm.forward_step( + curr_embed[:, 0, :], + torch.tensor([self.base_lm.kv_cache.step()], device=curr_embed.device), + ).clone() + lm_hidden = self.fsq_layer(lm_hidden) + residual_hidden = self.residual_lm.forward_step( + lm_hidden + curr_embed[:, 0, :], + torch.tensor([self.residual_lm.kv_cache.step()], device=curr_embed.device), + ).clone() + + if not streaming: + pred_feat_seq_cat = torch.cat(pred_feat_seq, dim=1) + feat_pred = rearrange(pred_feat_seq_cat, "b t p d -> b d (t p)", b=B, p=self.patch_size) + yield feat_pred, pred_feat_seq_cat.squeeze(0).cpu() + + @torch.inference_mode() + def generate_latents_with_prompt_cache( + self, + target_text: str, + prompt_cache: dict, + min_len: int = 2, + max_len: int = 2000, + inference_timesteps: int = 10, + cfg_value: float = 2.0, + retry_badcase: bool = False, + retry_badcase_max_times: int = 3, + retry_badcase_ratio_threshold: float = 6.0, + streaming_prefix_len: int = 3, + ) -> tuple[None, torch.Tensor, torch.Tensor]: + return next( + self._generate_with_prompt_cache( + target_text=target_text, + prompt_cache=prompt_cache, + min_len=min_len, + max_len=max_len, + inference_timesteps=inference_timesteps, + cfg_value=cfg_value, + retry_badcase=retry_badcase, + retry_badcase_max_times=retry_badcase_max_times, + retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, + streaming=False, + streaming_prefix_len=streaming_prefix_len, + latents_only=True, + ) + ) + + @torch.inference_mode() + def generate_latents_with_prompt_cache_streaming( + self, + target_text: str, + prompt_cache: dict, + min_len: int = 2, + max_len: int = 2000, + inference_timesteps: int = 10, + cfg_value: float = 2.0, + retry_badcase: bool = False, + retry_badcase_max_times: int = 3, + retry_badcase_ratio_threshold: float = 6.0, + streaming_prefix_len: int = 3, + ) -> Generator[tuple[None, torch.Tensor, torch.Tensor], None, None]: + return self._generate_with_prompt_cache( + target_text=target_text, + prompt_cache=prompt_cache, + min_len=min_len, + max_len=max_len, + inference_timesteps=inference_timesteps, + cfg_value=cfg_value, + retry_badcase=retry_badcase, + retry_badcase_max_times=retry_badcase_max_times, + retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, + streaming=True, + streaming_prefix_len=streaming_prefix_len, + latents_only=True, + ) + + @torch.inference_mode() + def _generate_with_prompt_cache( + self, + target_text: str, + prompt_cache: dict, + min_len: int = 2, + max_len: int = 2000, + inference_timesteps: int = 10, + cfg_value: float = 2.0, + retry_badcase: bool = False, + retry_badcase_max_times: int = 3, + retry_badcase_ratio_threshold: float = 6.0, + streaming: bool = False, + streaming_prefix_len: int = 3, + latents_only: bool = False, + ) -> Generator[tuple[torch.Tensor | None, torch.Tensor, torch.Tensor | list[torch.Tensor]], None, None]: + if retry_badcase and streaming: + warnings.warn("Retry on bad cases is not supported in streaming mode, setting retry_badcase=False.") + retry_badcase = False + if prompt_cache is None: + prompt_audio_feat = torch.empty((0, self.patch_size, self.audio_vae.latent_dim), dtype=torch.float32) + text = target_text + else: + prompt_audio_feat = prompt_cache["audio_feat"] + prompt_text = prompt_cache["prompt_text"] + text = prompt_text + target_text + + text_token = torch.LongTensor(self.text_tokenizer(text)) + text_token = torch.cat( + [ + text_token, + torch.tensor([self.audio_start_token], dtype=torch.int32, device=text_token.device), + ], + dim=-1, + ) + target_text_token = torch.LongTensor(self.text_tokenizer(target_text)) + + audio_length = prompt_audio_feat.size(0) + text_length = text_token.shape[0] + text_pad_token = torch.zeros(audio_length, dtype=torch.int32, device=text_token.device) + audio_pad_feat = torch.zeros( + (text_token.shape[0], self.patch_size, self.audio_vae.latent_dim), + dtype=torch.float32, + device=text_token.device, + ) + text_token = torch.cat([text_token, text_pad_token]) + audio_feat = torch.cat([audio_pad_feat, prompt_audio_feat], dim=0) + text_mask = ( + torch.cat([torch.ones(text_length), torch.zeros(audio_length)]).type(torch.int32).to(text_token.device) + ) + audio_mask = ( + torch.cat([torch.zeros(text_length), torch.ones(audio_length)]).type(torch.int32).to(text_token.device) + ) + + text_token = text_token.unsqueeze(0).to(self.device) + text_mask = text_mask.unsqueeze(0).to(self.device) + audio_feat = audio_feat.unsqueeze(0).to(self.device).to(get_dtype(self.config.dtype)) + audio_mask = audio_mask.unsqueeze(0).to(self.device) + + target_text_length = len(self.text_tokenizer(target_text)) + retry_badcase_times = 0 + while retry_badcase_times < retry_badcase_max_times: + inference_result = self._inference( + text_token, + text_mask, + audio_feat, + audio_mask, + min_len=min_len, + max_len=min(int(target_text_length * retry_badcase_ratio_threshold + 10), max_len), + inference_timesteps=inference_timesteps, + cfg_value=cfg_value, + streaming=streaming, + streaming_prefix_len=streaming_prefix_len, + ) + if streaming: + patch_len = self.patch_size * self.chunk_size + for latent_pred, pred_audio_feat in inference_result: + if latents_only: + decode_audio = None + yield (decode_audio, target_text_token, latent_pred) + else: + decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)) + decode_audio = decode_audio[..., -patch_len:].squeeze(1).cpu() + yield (decode_audio, target_text_token, pred_audio_feat) + break + + latent_pred, pred_audio_feat = next(inference_result) + if retry_badcase and pred_audio_feat.shape[0] >= target_text_length * retry_badcase_ratio_threshold: + ratio = pred_audio_feat.shape[0] / target_text_length + print(f" Badcase detected, audio_text_ratio={ratio}, retrying...", file=sys.stderr) + retry_badcase_times += 1 + continue + break + + if not streaming: + if latents_only: + decode_audio = None + else: + decode_audio = self.audio_vae.decode(latent_pred.to(torch.float32)) + patch_len = self.patch_size * self.chunk_size + if audio_mask.sum().item() > 0: + decode_audio = decode_audio[..., patch_len * (streaming_prefix_len - 1) :].squeeze(1).cpu() + else: + decode_audio = decode_audio[..., :].squeeze(1).cpu() + yield (decode_audio, target_text_token, pred_audio_feat) + + VoxCPMModelForOmni.__name__ = "VoxCPMModelForOmni" + VoxCPMModelForOmni.__qualname__ = "VoxCPMModelForOmni" + return VoxCPMModelForOmni + + +def _import_voxcpm_model_class() -> type[Any]: + base = _import_voxcpm_base_model_class() + return _make_voxcpm_model_for_omni(base) + + +def _load_native_voxcpm_model( + model_path: str, + *, + device: torch.device, + dtype: str | None, +): + VoxCPMModel = _import_voxcpm_model_class() + model_dir = resolve_voxcpm_model_dir(model_path) + runtime_model_path = _prepare_runtime_model_dir(model_dir, target_device=device, target_dtype=dtype) + + if device.type == "npu" and hasattr(torch, "npu"): + torch.npu.set_device(device) + + with _force_cuda_available_for_npu(device): + return VoxCPMModel.from_local( + runtime_model_path, + optimize=device.type == "cuda", + ) + + +def _load_native_voxcpm_latent_generator( + model_path: str, + *, + device: torch.device, + dtype: str | None, +) -> _DirectVoxCPMLatentGenerator: + return _DirectVoxCPMLatentGenerator(_load_native_voxcpm_model(model_path, device=device, dtype=dtype)) + + +def _load_native_voxcpm_audio_vae( + model_path: str, + *, + device: torch.device, +) -> _DirectVoxCPMAudioVAE: + AudioVAE, AudioVAEConfig = _import_voxcpm_audio_vae_classes() + model_dir = resolve_voxcpm_model_dir(model_path) + runtime_model_path = _prepare_runtime_model_dir(model_dir, target_device=device, target_dtype="float32") + config_dict = json.loads((Path(runtime_model_path) / "config.json").read_text()) + audio_vae_config = config_dict.get("audio_vae_config") + audio_vae = AudioVAE(config=AudioVAEConfig(**audio_vae_config)) if audio_vae_config is not None else AudioVAE() + + state_dict = torch.load( + Path(runtime_model_path) / "audiovae.pth", + map_location="cpu", + weights_only=True, + )["state_dict"] + audio_vae.load_state_dict(state_dict, strict=True) + audio_vae = audio_vae.to(device=device, dtype=torch.float32).eval() + if device.type == "npu" and hasattr(torch, "npu"): + torch.npu.set_device(device) + patch_size = int(config_dict.get("patch_size", 2)) + return _DirectVoxCPMAudioVAE(audio_vae, patch_size=patch_size) + + +class VoxCPMForConditionalGeneration(nn.Module): + input_modalities = "audio" + _LATENT_STAGES = {"latent_generator", "latent", "ar_dit"} + _VAE_STAGES = {"vae", "audio_vae"} + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + del prefix + self.vllm_config = vllm_config + self.model_path = vllm_config.model_config.model + self.model_stage = getattr(vllm_config.model_config, "model_stage", "latent_generator") + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + self.enable_update_additional_information = True + self.requires_raw_input_tokens = True + self.inject_omni_request_id_into_runtime_info = True + self._pipeline = None + self._latent_stream_gens: dict[str, Any] = {} + self._latent_stream_terminal_pending: dict[str, int] = {} + self._latent_stream_completed: set[str] = set() + self._next_local_stream_key = 0 + self._ar_emit_stop_token = True + + def _runner_hidden_device_dtype(self) -> tuple[torch.device, torch.dtype]: + device = _resolve_runtime_device(self.vllm_config) + model_config = getattr(self.vllm_config, "model_config", None) + dtype = getattr(model_config, "dtype", torch.float32) if model_config is not None else torch.float32 + return device, dtype + + def _ensure_model_loaded(self): + if self._pipeline is not None: + return + + target_device = _resolve_runtime_device(self.vllm_config) + model_dtype = getattr(self.vllm_config.model_config, "dtype", None) + normalized_dtype = _normalize_dtype_name(model_dtype) + if self.model_stage in self._LATENT_STAGES: + self._pipeline = _load_native_voxcpm_latent_generator( + self.model_path, + device=target_device, + dtype=normalized_dtype, + ) + elif self.model_stage in self._VAE_STAGES: + self._pipeline = _load_native_voxcpm_audio_vae( + self.model_path, + device=target_device, + ) + else: + raise ValueError( + f"Unsupported VoxCPM model_stage: {self.model_stage}. " + "pure_voxcpm only supports split-stage latent_generator/vae inference." + ) + + logger.info("Loaded VoxCPM stage '%s' on %s", self.model_stage, _device_to_string(target_device)) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + del weights + self._ensure_model_loaded() + return set() + + @staticmethod + def _extract_val(info: dict[str, Any], key: str, default: Any) -> Any: + value = info.get(key, default) + if isinstance(value, list): + return value[0] if value else default + return value + + def _resolve_stream_request_key(self, info: dict[str, Any]) -> str: + request_key = info.get("__voxcpm_stream_key") + if request_key is not None: + return str(request_key) + + request_key = info.get("_omni_req_id") + if request_key is not None: + request_key = str(request_key) + info["__voxcpm_stream_key"] = request_key + return request_key + + request_key = f"voxcpm-local-{self._next_local_stream_key}" + self._next_local_stream_key += 1 + info["__voxcpm_stream_key"] = request_key + return str(request_key) + + def _recover_latent_from_input_ids(self, input_ids: torch.Tensor | None) -> torch.Tensor | None: + if input_ids is None or input_ids.numel() == 0: + return None + flat_ids = input_ids.detach().reshape(-1).to("cpu") + if flat_ids.numel() < 4 or int(flat_ids[0].item()) != _VOXCPM_LATENT_MAGIC: + return None + latent_dim = int(flat_ids[1].item()) + time_dim = int(flat_ids[2].item()) + payload = flat_ids[3:] + expected = latent_dim * time_dim + if latent_dim <= 0 or time_dim <= 0: + raise ValueError(f"Invalid VoxCPM latent header: latent_dim={latent_dim}, time_dim={time_dim}") + if int(payload.numel()) != expected: + raise ValueError( + "Invalid VoxCPM latent payload size: " + f"expected={expected}, actual={int(payload.numel())}, " + f"latent_dim={latent_dim}, time_dim={time_dim}" + ) + packed = payload.to(dtype=torch.int32).to(torch.uint16) + return packed.view(torch.bfloat16).to(torch.float32).reshape(1, latent_dim, time_dim) + + def _maybe_recover_vae_infos( + self, + infos: list[dict[str, Any]], + input_ids: torch.Tensor | None, + *, + async_chunk: bool, + ) -> list[dict[str, Any]]: + if not async_chunk: + return infos + if any(self._extract_val(info, "latent_audio_feat", None) is not None for info in infos): + return infos + recovered = self._recover_latent_from_input_ids(input_ids) + if recovered is None: + return infos + return [{"latent_audio_feat": recovered}] + + @staticmethod + def _normalize_audio_samples(samples: Any) -> np.ndarray: + if isinstance(samples, torch.Tensor): + return samples.detach().cpu().float().reshape(-1).numpy() + return np.asarray(samples, dtype=np.float32).reshape(-1) + + @classmethod + def _normalize_ref_audio(cls, ref_audio: Any) -> tuple[np.ndarray, int]: + if isinstance(ref_audio, str): + raise TypeError("String ref_audio should be handled as a path before waveform normalization.") + + if isinstance(ref_audio, dict): + sample_rate = ref_audio.get("sample_rate") or ref_audio.get("sampling_rate") or ref_audio.get("sr") + samples = None + for key in ("audio", "wav", "samples", "array", "waveform"): + if key in ref_audio and ref_audio[key] is not None: + samples = ref_audio[key] + break + if sample_rate is None or samples is None: + raise ValueError("ref_audio dict must contain waveform data and sample rate.") + return cls._normalize_audio_samples(samples), int(sample_rate) + + if isinstance(ref_audio, (list, tuple)): + if len(ref_audio) == 1: + return cls._normalize_ref_audio(ref_audio[0]) + if len(ref_audio) == 2 and np.isscalar(ref_audio[1]): + return cls._normalize_audio_samples(ref_audio[0]), int(ref_audio[1]) + + raise TypeError(f"Unsupported ref_audio format: {type(ref_audio)!r}") + + @staticmethod + def _write_temp_prompt_wav(waveform: np.ndarray, sample_rate: int) -> str: + prompt_file = tempfile.NamedTemporaryFile(delete=False, suffix=".wav") + prompt_file.close() + + wav = np.asarray(waveform, dtype=np.float32).reshape(-1) + wav = np.clip(wav, -1.0, 1.0) + pcm16 = (wav * 32767.0).astype(np.int16) + with wave.open(prompt_file.name, "wb") as wav_file: + wav_file.setnchannels(1) + wav_file.setsampwidth(2) + wav_file.setframerate(int(sample_rate)) + wav_file.writeframes(pcm16.tobytes()) + + return prompt_file.name + + @classmethod + def _resolve_prompt_inputs(cls, info: dict[str, Any]) -> tuple[str | None, str | None, str | None]: + prompt_text = cls._extract_val(info, "prompt_text", None) + prompt_wav_path = cls._extract_val(info, "prompt_wav_path", None) + if prompt_wav_path: + if prompt_text is None: + prompt_text = cls._extract_val(info, "ref_text", None) + return prompt_wav_path, prompt_text, None + + ref_audio = cls._extract_val(info, "ref_audio", None) + ref_text = cls._extract_val(info, "ref_text", None) + if ref_audio is None or ref_text is None: + return None, None, None + if isinstance(ref_audio, str): + return ref_audio, ref_text, None + + waveform, sample_rate = cls._normalize_ref_audio(ref_audio) + temp_prompt_wav = cls._write_temp_prompt_wav(waveform, sample_rate) + return temp_prompt_wav, ref_text, temp_prompt_wav + + def embed_input_ids(self, input_ids: torch.Tensor, **_: Any) -> torch.Tensor: + if input_ids.numel() == 0: + return torch.empty((0, 1), device=input_ids.device, dtype=torch.float32) + return torch.zeros((input_ids.shape[0], 1), device=input_ids.device, dtype=torch.float32) + + def _get_vocab_size(self) -> int: + model_config = getattr(self.vllm_config, "model_config", None) + if model_config is not None: + getter = getattr(model_config, "get_vocab_size", None) + if callable(getter): + try: + return int(getter()) + except Exception: + pass + hf_config = getattr(model_config, "hf_text_config", None) + if hf_config is not None and hasattr(hf_config, "vocab_size"): + return int(hf_config.vocab_size) + return 32000 + + def _make_empty_output( + self, + *, + output_key: str, + payload_factory: Callable[[], torch.Tensor], + infos: list[dict[str, Any]], + sample_rate: int, + out_device: torch.device, + out_dtype: torch.dtype, + hidden_rows: int | None = None, + ) -> OmniOutput: + if hidden_rows is None: + hidden_rows = len(infos) + return OmniOutput( + text_hidden_states=torch.zeros((hidden_rows, 1), device=out_device, dtype=out_dtype), + multimodal_outputs={ + output_key: [payload_factory() for _ in infos], + "sr": [torch.tensor(sample_rate, dtype=torch.int32) for _ in infos], + }, + ) + + def _finalize_stage_output( + self, + *, + output_key: str, + outputs: list[torch.Tensor], + sample_rates: list[torch.Tensor], + out_device: torch.device, + out_dtype: torch.dtype, + hidden_rows: int | None = None, + ) -> OmniOutput: + multimodal_outputs: dict[str, Any] = {output_key: outputs, "sr": sample_rates} + if hidden_rows is not None: + text_hidden_states = torch.zeros((hidden_rows, 1), device=out_device, dtype=out_dtype) + elif outputs: + outputs_tensor = torch.stack(outputs) + text_hidden_states = ( + outputs_tensor.unsqueeze(-1) + if outputs_tensor.ndim == 1 + else outputs_tensor.reshape(-1, outputs_tensor.shape[-1]) + ) + else: + text_hidden_states = torch.zeros((0, 1), device=out_device, dtype=out_dtype) + text_hidden_states = text_hidden_states.to(device=out_device, dtype=out_dtype) + return OmniOutput( + text_hidden_states=text_hidden_states, + multimodal_outputs=multimodal_outputs, + ) + + def _forward_vae_stage( + self, + infos: list[dict[str, Any]], + *, + sample_rate: int, + async_chunk: bool, + out_device: torch.device, + out_dtype: torch.dtype, + ) -> OmniOutput: + if all(self._extract_val(info, "latent_audio_feat", None) is None for info in infos): + self._ar_emit_stop_token = True + return self._make_empty_output( + output_key="model_outputs", + payload_factory=lambda: torch.zeros((0,), dtype=torch.float32), + infos=infos, + sample_rate=sample_rate, + out_device=out_device, + out_dtype=out_dtype, + ) + + outputs: list[torch.Tensor] = [] + sample_rates: list[torch.Tensor] = [] + for info in infos: + latent_audio_feat = self._extract_val(info, "latent_audio_feat", None) + audio_tensor = self._pipeline.decode(latent_audio_feat, trim_streaming_patch=async_chunk) + outputs.append(audio_tensor.float().cpu()) + sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32)) + + self._ar_emit_stop_token = True + return self._finalize_stage_output( + output_key="model_outputs", + outputs=outputs, + sample_rates=sample_rates, + out_device=out_device, + out_dtype=out_dtype, + ) + + def _forward_latent_stage( + self, + infos: list[dict[str, Any]], + *, + sample_rate: int, + async_chunk: bool, + out_device: torch.device, + out_dtype: torch.dtype, + hidden_rows: int, + ) -> OmniOutput: + texts = [self._extract_val(info, "text", "") for info in infos] + if all(not text for text in texts): + self._ar_emit_stop_token = True + return self._make_empty_output( + output_key="latent_audio_feat", + payload_factory=lambda: torch.zeros((0,), dtype=torch.float32), + infos=infos, + sample_rate=sample_rate, + out_device=out_device, + out_dtype=out_dtype, + hidden_rows=hidden_rows, + ) + + outputs: list[torch.Tensor] = [] + sample_rates: list[torch.Tensor] = [] + last_chunk_flags: list[bool] | None = [] if async_chunk else None + payload_finished_flags: list[bool] | None = [] if async_chunk else None + for info in infos: + text = self._extract_val(info, "text", "") + cfg_value = float(self._extract_val(info, "cfg_value", 2.0)) + inference_timesteps = int(self._extract_val(info, "inference_timesteps", 10)) + min_len = int(self._extract_val(info, "min_len", 2)) + max_len = int(self._extract_val(info, "max_len", self._extract_val(info, "max_new_tokens", 4096))) + retry_badcase = bool(self._extract_val(info, "retry_badcase", True)) + retry_badcase_max_times = int(self._extract_val(info, "retry_badcase_max_times", 3)) + retry_badcase_ratio_threshold = float(self._extract_val(info, "retry_badcase_ratio_threshold", 6.0)) + streaming_prefix_len = int(self._extract_val(info, "streaming_prefix_len", 3)) + + request_key = self._resolve_stream_request_key(info) + created_temp: str | None = None + + if async_chunk: + terminal_pending = self._latent_stream_terminal_pending.get(request_key, 0) + if terminal_pending > 0: + outputs.append(torch.zeros((0,), dtype=torch.float32)) + assert last_chunk_flags is not None + last_chunk_flags.append(True) + assert payload_finished_flags is not None + payload_finished_flags.append(terminal_pending == 1) + if terminal_pending == 1: + self._latent_stream_terminal_pending.pop(request_key, None) + else: + self._latent_stream_terminal_pending[request_key] = terminal_pending - 1 + sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32)) + continue + + if request_key in self._latent_stream_completed: + outputs.append(torch.zeros((0,), dtype=torch.float32)) + assert last_chunk_flags is not None + last_chunk_flags.append(True) + assert payload_finished_flags is not None + payload_finished_flags.append(False) + sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32)) + continue + + if request_key not in self._latent_stream_gens: + prompt_wav_path, prompt_text, temp_prompt_wav = self._resolve_prompt_inputs(info) + created_temp = temp_prompt_wav + self._latent_stream_gens[request_key] = self._pipeline.iter_latent_chunks_streaming( + text=text, + prompt_wav_path=prompt_wav_path, + prompt_text=prompt_text, + cfg_value=cfg_value, + inference_timesteps=inference_timesteps, + min_len=min_len, + max_len=max_len, + streaming_prefix_len=streaming_prefix_len, + retry_badcase=False, + retry_badcase_max_times=retry_badcase_max_times, + retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, + ) + generator = self._latent_stream_gens[request_key] + try: + chunk_latent, is_last = next(generator) + except StopIteration: + self._latent_stream_gens.pop(request_key, None) + self._latent_stream_terminal_pending[request_key] = 1 + self._latent_stream_completed.add(request_key) + outputs.append(torch.zeros((0,), dtype=torch.float32)) + assert last_chunk_flags is not None + last_chunk_flags.append(True) + assert payload_finished_flags is not None + payload_finished_flags.append(True) + else: + if is_last: + self._latent_stream_gens.pop(request_key, None) + self._latent_stream_terminal_pending[request_key] = 1 + self._latent_stream_completed.add(request_key) + outputs.append(chunk_latent.detach().float().cpu()) + assert last_chunk_flags is not None + last_chunk_flags.append(bool(is_last)) + assert payload_finished_flags is not None + payload_finished_flags.append(False) + finally: + if created_temp is not None and os.path.exists(created_temp): + os.unlink(created_temp) + sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32)) + continue + + prompt_wav_path, prompt_text, temp_prompt_wav = self._resolve_prompt_inputs(info) + try: + latent_audio_feat = self._pipeline.generate_latents( + text=text, + prompt_wav_path=prompt_wav_path, + prompt_text=prompt_text, + cfg_value=cfg_value, + inference_timesteps=inference_timesteps, + min_len=min_len, + max_len=max_len, + retry_badcase=retry_badcase, + retry_badcase_max_times=retry_badcase_max_times, + retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, + ) + outputs.append(latent_audio_feat.float().cpu()) + finally: + if temp_prompt_wav is not None and os.path.exists(temp_prompt_wav): + os.unlink(temp_prompt_wav) + + sample_rates.append(torch.tensor(sample_rate, dtype=torch.int32)) + + self._ar_emit_stop_token = all(last_chunk_flags) if async_chunk and last_chunk_flags else True + output = self._finalize_stage_output( + output_key="latent_audio_feat", + outputs=outputs, + sample_rates=sample_rates, + out_device=out_device, + out_dtype=out_dtype, + hidden_rows=hidden_rows, + ) + if async_chunk and payload_finished_flags is not None: + output.multimodal_outputs["finished"] = [ + torch.tensor(flag, dtype=torch.bool) for flag in payload_finished_flags + ] + return output + + def compute_logits(self, hidden_states: torch.Tensor | OmniOutput, sampling_metadata: Any = None) -> torch.Tensor: + del sampling_metadata + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + if hidden_states is None: + device, dtype = self._runner_hidden_device_dtype() + hidden_states = torch.zeros((0, 1), device=device, dtype=dtype) + if hidden_states.ndim == 1: + hidden_states = hidden_states.unsqueeze(-1) + elif hidden_states.ndim > 2: + hidden_states = hidden_states.reshape(-1, hidden_states.shape[-1]) + + vocab_size = self._get_vocab_size() + num_rows = int(hidden_states.shape[0]) + logits = torch.zeros((num_rows, vocab_size), dtype=torch.float32, device=hidden_states.device) + eos_id = 2 if vocab_size > 2 else 0 + safe_id = 1 if vocab_size > 1 and 1 != eos_id else 0 + emit_stop = getattr(self, "_ar_emit_stop_token", True) + if num_rows > 0: + if emit_stop: + logits[:, eos_id] = 1.0e6 + else: + logits[:, eos_id] = -1.0e9 + logits[:, safe_id] = 1.0e6 + return logits + + @torch.no_grad() + def forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: Any = None, + inputs_embeds: torch.Tensor | None = None, + runtime_additional_information: list[dict[str, Any]] | None = None, + model_intermediate_buffer: list[dict[str, Any]] | None = None, + **kwargs: Any, + ) -> OmniOutput: + del positions, intermediate_tensors, inputs_embeds, kwargs + self._ensure_model_loaded() + out_device, out_dtype = self._runner_hidden_device_dtype() + if input_ids is not None and input_ids.device.type == out_device.type: + out_device = input_ids.device + + infos = model_intermediate_buffer or runtime_additional_information or [{}] + hidden_rows = len(infos) + if input_ids is not None and len(input_ids.shape) > 0: + hidden_rows = max(hidden_rows, int(input_ids.shape[0])) + sample_rate = int(getattr(self._pipeline, "sample_rate", 24000)) + async_chunk = bool(getattr(self.vllm_config.model_config, "async_chunk", False)) + if self.model_stage in self._VAE_STAGES: + infos = self._maybe_recover_vae_infos(infos, input_ids, async_chunk=async_chunk) + return self._forward_vae_stage( + infos, + sample_rate=sample_rate, + async_chunk=async_chunk, + out_device=out_device, + out_dtype=out_dtype, + ) + if self.model_stage in self._LATENT_STAGES: + return self._forward_latent_stage( + infos, + sample_rate=sample_rate, + async_chunk=async_chunk, + out_device=out_device, + out_dtype=out_dtype, + hidden_rows=hidden_rows, + ) + raise ValueError(f"Unsupported VoxCPM model_stage at runtime: {self.model_stage}") + + def make_empty_intermediate_tensors( + self, batch_size: int, dtype: torch.dtype, device: torch.device + ) -> IntermediateTensors: + del batch_size, dtype, device + return {} + + +__all__ = ["VoxCPMForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/voxcpm/voxcpm_loader.py b/vllm_omni/model_executor/models/voxcpm/voxcpm_loader.py new file mode 100644 index 000000000..dac7117ca --- /dev/null +++ b/vllm_omni/model_executor/models/voxcpm/voxcpm_loader.py @@ -0,0 +1,247 @@ +from __future__ import annotations + +import importlib +import json +import os +import shutil +import sys +import tempfile +from contextlib import contextmanager +from hashlib import sha256 +from pathlib import Path +from typing import Any +from unittest.mock import patch + +import numpy as np +import torch +from vllm.config import VllmConfig +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def _iter_voxcpm_src_candidates() -> list[Path]: + candidates: list[Path] = [] + env_path = os.environ.get("VLLM_OMNI_VOXCPM_CODE_PATH") + if env_path: + candidates.append(Path(env_path).expanduser()) + + repo_root = Path(__file__).resolve().parents[4] + candidates.append(repo_root.parent / "VoxCPM" / "src") + + unique_candidates: list[Path] = [] + seen: set[str] = set() + for candidate in candidates: + candidate_key = str(candidate) + if candidate_key in seen: + continue + seen.add(candidate_key) + unique_candidates.append(candidate) + return unique_candidates + + +def _prepend_voxcpm_src(candidate: Path) -> None: + candidate_str = str(candidate) + if candidate_str not in sys.path: + sys.path.insert(0, candidate_str) + + +def _import_voxcpm_attrs(module_name: str, *attr_names: str) -> tuple[Any, ...]: + last_exc: ImportError | None = None + for candidate in _iter_voxcpm_src_candidates(): + if not candidate.exists(): + continue + _prepend_voxcpm_src(candidate) + try: + module = importlib.import_module(module_name) + return tuple(getattr(module, attr_name) for attr_name in attr_names) + except ImportError as exc: + last_exc = exc + + try: + module = importlib.import_module(module_name) + return tuple(getattr(module, attr_name) for attr_name in attr_names) + except ImportError as exc: + last_exc = exc + + raise ImportError(f"Failed to import {module_name}.") from last_exc + + +def _import_voxcpm_base_model_class(): + """Import upstream ``VoxCPMModel`` from ``VoxCPM/src/voxcpm`` (env, sibling tree, or pip).""" + try: + (VoxCPMModel,) = _import_voxcpm_attrs("voxcpm.model.voxcpm", "VoxCPMModel") + return VoxCPMModel + except ImportError as exc: + raise ImportError( + "Failed to import VoxCPMModel. Install the `voxcpm` package or set " + "`VLLM_OMNI_VOXCPM_CODE_PATH` to the VoxCPM repository `src` directory " + "(the parent of the `voxcpm` package that contains `model/` and `modules/`)." + ) from exc + + +def _import_voxcpm_audio_vae_classes(): + try: + return _import_voxcpm_attrs("voxcpm.modules.audiovae", "AudioVAE", "AudioVAEConfig") + except ImportError as exc: + raise ImportError( + "Failed to import VoxCPM AudioVAE. Install the `voxcpm` package or set " + "`VLLM_OMNI_VOXCPM_CODE_PATH` to the VoxCPM repository `src` directory." + ) from exc + + +def _device_to_string(device: torch.device) -> str: + if device.index is None: + return device.type + return f"{device.type}:{device.index}" + + +def _normalize_dtype_name(dtype: Any) -> str | None: + if dtype is None: + return None + if isinstance(dtype, torch.dtype): + mapping = { + torch.bfloat16: "bfloat16", + torch.float16: "float16", + torch.float32: "float32", + } + return mapping.get(dtype, str(dtype).removeprefix("torch.")) + dtype_str = str(dtype) + return dtype_str.removeprefix("torch.") + + +def _resolve_runtime_device(vllm_config: VllmConfig) -> torch.device: + try: + from vllm_omni.platforms import current_omni_platform + + return current_omni_platform.get_torch_device() + except Exception: + pass + + device = getattr(getattr(vllm_config, "device_config", None), "device", None) + if isinstance(device, torch.device): + return device + if device: + return torch.device(device) + return torch.device("cpu") + + +def _prepare_runtime_model_dir( + model_path: str | Path, + *, + target_device: torch.device, + target_dtype: str | None, +) -> str: + source_dir = Path(model_path) + config_path = source_dir / "config.json" + if not config_path.exists(): + return str(source_dir) + + config_text = config_path.read_text() + config_dict = json.loads(config_text) + desired_device = target_device.type + desired_dtype = target_dtype or config_dict.get("dtype") + + if config_dict.get("device") == desired_device and config_dict.get("dtype") == desired_dtype: + return str(source_dir) + + digest = sha256(f"{source_dir.resolve()}:{config_text}:{desired_device}:{desired_dtype}".encode()).hexdigest()[:16] + runtime_dir = Path(tempfile.gettempdir()) / "vllm_omni_voxcpm_runtime" / digest + runtime_dir.mkdir(parents=True, exist_ok=True) + + for entry in source_dir.iterdir(): + target = runtime_dir / entry.name + if entry.name == "config.json" or target.exists(): + continue + try: + target.symlink_to(entry, target_is_directory=entry.is_dir()) + except OSError as exc: + logger.warning( + "Falling back to copying VoxCPM runtime artifact %s into %s because symlink creation failed: %s", + entry, + runtime_dir, + exc, + ) + if entry.is_dir(): + shutil.copytree(entry, target, dirs_exist_ok=True) + else: + shutil.copy2(entry, target) + + patched_config = dict(config_dict) + patched_config["device"] = desired_device + if desired_dtype is not None: + patched_config["dtype"] = desired_dtype + (runtime_dir / "config.json").write_text(json.dumps(patched_config, indent=2, sort_keys=True)) + return str(runtime_dir) + + +@contextmanager +def _force_cuda_available_for_npu(device: torch.device): + if device.type != "npu": + yield + return + + with patch("torch.cuda.is_available", return_value=True): + yield + + +def _is_torchcodec_load_error(exc: BaseException) -> bool: + message = str(exc).lower() + return "torchcodec" in message or "load_with_torchcodec" in message + + +def _load_audio_with_soundfile( + prompt_wav_path: str, + *, + sample_rate: int, +) -> torch.Tensor: + try: + import soundfile as sf + except ImportError: + raise + + audio_np, source_sr = sf.read(prompt_wav_path, dtype="float32", always_2d=True) + audio = torch.from_numpy(np.ascontiguousarray(audio_np.T)) + + if audio.size(0) > 1: + audio = audio.mean(dim=0, keepdim=True) + + if int(source_sr) != int(sample_rate): + try: + import torchaudio + except ImportError as exc: + raise ImportError("torchaudio is required for resampling prompt audio.") from exc + audio = torchaudio.functional.resample(audio, int(source_sr), int(sample_rate)) + + return audio + + +def _build_prompt_cache_with_soundfile(model: Any, *args: Any, **kwargs: Any) -> dict[str, Any]: + if args: + prompt_text = args[0] + prompt_wav_path = args[1] if len(args) > 1 else kwargs.get("prompt_wav_path") + else: + prompt_text = kwargs.get("prompt_text") + prompt_wav_path = kwargs.get("prompt_wav_path") + + if not prompt_text or not prompt_wav_path: + raise ValueError("prompt_text and prompt_wav_path are required") + + audio = _load_audio_with_soundfile(prompt_wav_path, sample_rate=int(model.sample_rate)) + + patch_len = model.patch_size * model.chunk_size + if audio.size(1) % patch_len != 0: + padding_size = patch_len - audio.size(1) % patch_len + audio = torch.nn.functional.pad(audio, (padding_size, 0)) + + audio_feat = model.audio_vae.encode(audio.to(model.device), model.sample_rate).cpu() + audio_feat = audio_feat.view( + model.audio_vae.latent_dim, + -1, + model.patch_size, + ).permute(1, 2, 0) + + return { + "prompt_text": prompt_text, + "audio_feat": audio_feat, + } diff --git a/vllm_omni/model_executor/models/voxcpm/voxcpm_runtime_utils.py b/vllm_omni/model_executor/models/voxcpm/voxcpm_runtime_utils.py new file mode 100644 index 000000000..36b4282c2 --- /dev/null +++ b/vllm_omni/model_executor/models/voxcpm/voxcpm_runtime_utils.py @@ -0,0 +1,44 @@ +from __future__ import annotations + +import json +import shutil +from pathlib import Path + + +def resolve_voxcpm_model_dir(model: str) -> Path: + model_path = Path(model).expanduser() + if model_path.exists(): + return model_path + + from huggingface_hub import snapshot_download + + return Path(snapshot_download(repo_id=model)) + + +def prepare_voxcpm_hf_config_dir(model_dir: str | Path, hf_config_dir: str | Path) -> Path: + model_dir = Path(model_dir).expanduser() + hf_config_dir = Path(hf_config_dir).expanduser() + hf_config_dir.mkdir(parents=True, exist_ok=True) + + source_config_path = model_dir / "config.json" + if not source_config_path.exists(): + raise FileNotFoundError(f"VoxCPM config.json not found under {model_dir}") + + config_path = hf_config_dir / "config.json" + shutil.copy2(source_config_path, config_path) + + source_generation_config_path = model_dir / "generation_config.json" + if source_generation_config_path.exists(): + shutil.copy2(source_generation_config_path, hf_config_dir / "generation_config.json") + + config_dict = json.loads(config_path.read_text(encoding="utf-8")) + config_dict["model_type"] = "voxcpm" + config_dict.setdefault("architectures", ["VoxCPMForConditionalGeneration"]) + config_path.write_text(json.dumps(config_dict, indent=2, ensure_ascii=False), encoding="utf-8") + return hf_config_dir + + +__all__ = [ + "prepare_voxcpm_hf_config_dir", + "resolve_voxcpm_model_dir", +] diff --git a/vllm_omni/model_executor/models/voxcpm/voxcpm_stage_wrappers.py b/vllm_omni/model_executor/models/voxcpm/voxcpm_stage_wrappers.py new file mode 100644 index 000000000..f4446c796 --- /dev/null +++ b/vllm_omni/model_executor/models/voxcpm/voxcpm_stage_wrappers.py @@ -0,0 +1,185 @@ +from __future__ import annotations + +import os +from collections.abc import Generator +from typing import Any + +import torch +import torch.nn as nn +from einops import rearrange + + +class _DirectVoxCPMLatentGenerator: + def __init__(self, tts_model: Any): + self.tts_model = tts_model + self.sample_rate = int(getattr(tts_model, "sample_rate", 24000)) + + def generate_latents( + self, + *, + text: str, + prompt_wav_path: str | None = None, + prompt_text: str | None = None, + cfg_value: float = 2.0, + inference_timesteps: int = 10, + min_len: int = 2, + max_len: int = 4096, + retry_badcase: bool = True, + retry_badcase_max_times: int = 3, + retry_badcase_ratio_threshold: float = 6.0, + ) -> torch.Tensor: + if not isinstance(text, str) or not text.strip(): + raise ValueError("target text must be a non-empty string") + if (prompt_wav_path is None) != (prompt_text is None): + raise ValueError("prompt_wav_path and prompt_text must both be provided or both be None") + if prompt_wav_path is not None and not os.path.exists(prompt_wav_path): + raise FileNotFoundError(f"prompt_wav_path does not exist: {prompt_wav_path}") + + prompt_cache = None + if prompt_wav_path is not None and prompt_text is not None: + prompt_cache = self.tts_model.build_prompt_cache( + prompt_text=prompt_text, + prompt_wav_path=prompt_wav_path, + ) + + gen_kw = dict( + target_text=" ".join(text.split()), + prompt_cache=prompt_cache, + min_len=min_len, + max_len=max_len, + inference_timesteps=inference_timesteps, + cfg_value=cfg_value, + retry_badcase=retry_badcase, + retry_badcase_max_times=retry_badcase_max_times, + retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, + ) + latent_entry = getattr(self.tts_model, "generate_latents_with_prompt_cache", None) + if latent_entry is not None: + _, _, pred_audio_feat = latent_entry(**gen_kw) + else: + try: + _, _, pred_audio_feat = self.tts_model.generate_with_prompt_cache( + **gen_kw, + latents_only=True, + ) + except TypeError: + _, _, pred_audio_feat = self.tts_model.generate_with_prompt_cache(**gen_kw) + return pred_audio_feat.detach().cpu().to(torch.float32) + + def iter_latent_chunks_streaming( + self, + *, + text: str, + prompt_wav_path: str | None = None, + prompt_text: str | None = None, + cfg_value: float = 2.0, + inference_timesteps: int = 10, + min_len: int = 2, + max_len: int = 4096, + streaming_prefix_len: int = 3, + retry_badcase: bool = False, + retry_badcase_max_times: int = 3, + retry_badcase_ratio_threshold: float = 6.0, + ) -> Generator[tuple[torch.Tensor, bool], None, None]: + """Yield ``(latent_window, is_last_chunk)`` for Omni async_chunk latent to VAE.""" + if not isinstance(text, str) or not text.strip(): + raise ValueError("target text must be a non-empty string") + if (prompt_wav_path is None) != (prompt_text is None): + raise ValueError("prompt_wav_path and prompt_text must both be provided or both be None") + if prompt_wav_path is not None and not os.path.exists(prompt_wav_path): + raise FileNotFoundError(f"prompt_wav_path does not exist: {prompt_wav_path}") + + prompt_cache = None + if prompt_wav_path is not None and prompt_text is not None: + prompt_cache = self.tts_model.build_prompt_cache( + prompt_text=prompt_text, + prompt_wav_path=prompt_wav_path, + ) + + gen_kw = dict( + target_text=" ".join(text.split()), + prompt_cache=prompt_cache, + min_len=min_len, + max_len=max_len, + inference_timesteps=inference_timesteps, + cfg_value=cfg_value, + retry_badcase=retry_badcase, + retry_badcase_max_times=retry_badcase_max_times, + retry_badcase_ratio_threshold=retry_badcase_ratio_threshold, + streaming_prefix_len=streaming_prefix_len, + ) + stream_entry = getattr(self.tts_model, "generate_latents_with_prompt_cache_streaming", None) + if stream_entry is not None: + gen = stream_entry(**gen_kw) + else: + fallback_stream_entry = getattr(self.tts_model, "generate_with_prompt_cache_streaming", None) + if fallback_stream_entry is not None: + gen = fallback_stream_entry(**gen_kw, latents_only=True) + else: + gen = self.tts_model._generate_with_prompt_cache(streaming=True, latents_only=True, **gen_kw) + + iterator = iter(gen) + previous = next(iterator, None) + while previous is not None: + current = next(iterator, None) + _, _target_tok, chunk_latent = previous + if not isinstance(chunk_latent, torch.Tensor): + chunk_latent = torch.as_tensor(chunk_latent) + yield chunk_latent, current is None + previous = current + + +class _DirectVoxCPMAudioVAE: + def __init__(self, audio_vae: nn.Module, *, patch_size: int = 2): + self.audio_vae = audio_vae + self.sample_rate = int(getattr(audio_vae, "sample_rate", 24000)) + self.latent_dim = int(getattr(audio_vae, "latent_dim", 64)) + self.patch_size = int(patch_size) + self._chunk_size = int(getattr(audio_vae, "chunk_size", 1)) + self._stream_audio_patch_samples = max(1, self.patch_size * self._chunk_size) + + def _prepare_latents_for_decode(self, latent_audio_feat: Any) -> torch.Tensor: + latents = latent_audio_feat + if not isinstance(latents, torch.Tensor): + latents = torch.tensor(latents, dtype=torch.float32) + latents = latents.detach().to(torch.float32) + + if latents.ndim == 3: + if latents.shape[-1] == self.latent_dim: + latents = rearrange(latents, "t p d -> 1 d (t p)") + elif latents.shape[1] == self.latent_dim: + latents = latents.contiguous() + else: + raise ValueError(f"Unsupported latent_audio_feat shape: {tuple(latents.shape)}") + elif latents.ndim == 2: + if latents.shape[0] == self.latent_dim: + latents = latents.unsqueeze(0) + elif latents.shape[1] == self.latent_dim: + latents = rearrange(latents, "t d -> 1 d t") + else: + raise ValueError(f"Unsupported latent_audio_feat shape: {tuple(latents.shape)}") + else: + raise ValueError(f"Unsupported latent_audio_feat ndim: {latents.ndim}") + + return latents + + @torch.no_grad() + def decode(self, latent_audio_feat: Any, *, trim_streaming_patch: bool = False) -> torch.Tensor: + latents = self._prepare_latents_for_decode(latent_audio_feat) + device = next(self.audio_vae.parameters()).device + raw = self.audio_vae.decode(latents.to(device=device, dtype=torch.float32)) + if isinstance(raw, dict): + audio = raw.get("audio") + if audio is None: + audio = next(v for v in raw.values() if isinstance(v, torch.Tensor)) + else: + audio = raw + if audio.dim() == 3: + stream = audio.squeeze(1) + elif audio.dim() == 2: + stream = audio + else: + stream = audio.reshape(audio.shape[0], -1) + if trim_streaming_patch: + stream = stream[..., -self._stream_audio_patch_samples :] + return stream.reshape(-1).detach().cpu().to(torch.float32) diff --git a/vllm_omni/model_executor/stage_configs/voxcpm.yaml b/vllm_omni/model_executor/stage_configs/voxcpm.yaml new file mode 100644 index 000000000..a5f324f66 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/voxcpm.yaml @@ -0,0 +1,69 @@ +# VoxCPM two-stage (latent → VAE) without async_chunk: one-shot latent then decode. +stage_args: + - stage_id: 0 + stage_type: llm + is_comprehension: true + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + dtype: bfloat16 + model_stage: latent_generator + model_arch: VoxCPMForConditionalGeneration + # Optional persistent HF-compatible config dir for native VoxCPM models. + hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: latent + gpu_memory_utilization: 0.7 + distributed_executor_backend: "mp" + max_num_batched_tokens: 4096 + max_model_len: 4096 + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 4096 + stop_token_ids: [2] + seed: 42 + detokenize: false + repetition_penalty: 1.0 + final_output: false + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + dtype: float32 + model_stage: vae + model_arch: VoxCPMForConditionalGeneration + hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio + gpu_memory_utilization: 0.15 + distributed_executor_backend: "mp" + max_num_batched_tokens: 8192 + max_model_len: 4096 + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.voxcpm.latent2vae + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 1 + seed: 42 + detokenize: true + repetition_penalty: 1.0 diff --git a/vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml b/vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml new file mode 100644 index 000000000..cf78d4e43 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/voxcpm_async_chunk.yaml @@ -0,0 +1,102 @@ +# VoxCPM two-stage streaming (align with qwen3_tts.yaml async_chunk pattern). +# Stage0 (latent_generator) emits latent in time chunks; Stage1 (VAE) decodes as chunks arrive. +async_chunk: true +stage_args: + - stage_id: 0 + stage_type: llm + is_comprehension: true + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + dtype: bfloat16 + model_stage: latent_generator + model_arch: VoxCPMForConditionalGeneration + # Optional persistent HF-compatible config dir for native VoxCPM models. + hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: true + enable_prefix_caching: false + engine_output_type: latent + gpu_memory_utilization: 0.7 + distributed_executor_backend: "mp" + max_num_batched_tokens: 4096 + max_model_len: 4096 + custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxcpm.latent2vae_async_chunk + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 4096 + stop_token_ids: [2] + seed: 42 + detokenize: false + repetition_penalty: 1.0 + final_output: false + output_connectors: + to_stage_1: voxcpm_shm + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + dtype: float32 + model_stage: vae + model_arch: VoxCPMForConditionalGeneration + hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio + gpu_memory_utilization: 0.15 + distributed_executor_backend: "mp" + max_num_batched_tokens: 8192 + max_model_len: 4096 + engine_input_source: [0] + final_output: true + final_output_type: audio + input_connectors: + from_stage_0: voxcpm_shm + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 128 + seed: 42 + detokenize: true + repetition_penalty: 1.0 + + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + + connectors: + voxcpm_shm: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + # Frame-aligned codec streaming transport. + codec_streaming: true + # Connector polling / timeout (unit: loop count, sleep interval in seconds). + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + # Align with Omni: small chunks with sufficient context overlap. + codec_chunk_frames: 1 + codec_left_context_frames: 1 + + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/vllm_omni/model_executor/stage_input_processors/voxcpm.py b/vllm_omni/model_executor/stage_input_processors/voxcpm.py new file mode 100644 index 000000000..c2fcf521b --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/voxcpm.py @@ -0,0 +1,128 @@ +from __future__ import annotations + +from typing import Any + +import torch +from vllm.inputs import TextPrompt + +from vllm_omni.inputs.data import OmniTokensPrompt + +_VOXCPM_LATENT_MAGIC = 131071 + + +def _serialize_latent_to_codes(latent: Any) -> list[int]: + latent_tensor = latent if isinstance(latent, torch.Tensor) else torch.as_tensor(latent) + latent_tensor = latent_tensor.detach().cpu().contiguous() + if latent_tensor.ndim == 3: + if latent_tensor.shape[0] != 1: + raise ValueError(f"Expected batch=1 latent tensor, got shape={tuple(latent_tensor.shape)}") + latent_tensor = latent_tensor.squeeze(0) + if latent_tensor.ndim != 2: + raise ValueError(f"Unsupported latent_audio_feat shape for async chunk: {tuple(latent_tensor.shape)}") + latent_dim, time_dim = int(latent_tensor.shape[0]), int(latent_tensor.shape[1]) + packed = latent_tensor.to(torch.bfloat16).contiguous().view(torch.uint16).reshape(-1).to(torch.int32) + return [_VOXCPM_LATENT_MAGIC, latent_dim, time_dim, *packed.tolist()] + + +def _coerce_finished_flag(value: Any) -> bool: + """Normalize VoxCPM async-chunk finished markers to a Python bool.""" + if value is None: + return False + if isinstance(value, torch.Tensor): + if value.numel() != 1: + raise ValueError(f"finished tensor must be scalar, got shape={tuple(value.shape)}") + return bool(value.detach().cpu().item()) + if isinstance(value, (list, tuple)): + if not value: + return False + if len(value) != 1: + raise ValueError(f"finished container must have one element, got len={len(value)}") + return _coerce_finished_flag(value[0]) + return bool(value) + + +def latent2vae( + stage_list: list[Any], + engine_input_source: list[int], + prompt: OmniTokensPrompt | TextPrompt | None = None, + requires_multimodal_data: bool = False, +) -> list[OmniTokensPrompt]: + del prompt, requires_multimodal_data + + if not engine_input_source: + raise ValueError("engine_input_source cannot be empty") + + source_stage_id = engine_input_source[0] + if source_stage_id >= len(stage_list): + raise IndexError(f"Invalid stage_id: {source_stage_id}") + + source_outputs = stage_list[source_stage_id].engine_outputs + if source_outputs is None: + raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") + + vae_inputs: list[OmniTokensPrompt] = [] + for source_output in source_outputs: + output = source_output.outputs[0] + multimodal_output = getattr(output, "multimodal_output", None) + if not isinstance(multimodal_output, dict) or "latent_audio_feat" not in multimodal_output: + raise ValueError( + "VoxCPM latent stage output missing 'latent_audio_feat'. " + f"request_id={getattr(source_output, 'request_id', None)}" + ) + + additional_information = { + "latent_audio_feat": multimodal_output["latent_audio_feat"], + } + if "sr" in multimodal_output: + additional_information["sample_rate"] = [int(multimodal_output["sr"])] + + vae_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0], + additional_information=additional_information, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + + return vae_inputs + + +def latent2vae_async_chunk( + transfer_manager: Any, + pooling_output: dict[str, Any] | None, + request: Any, + is_finished: bool = False, +) -> dict[str, Any] | None: + """Stage-0 latent → stage-1 VAE under ``async_chunk`` (connector payload).""" + # Kept for callback signature compatibility with OmniChunkTransferAdapter. + _ = transfer_manager + finished_request = _coerce_finished_flag(is_finished) + if callable(getattr(request, "is_finished", None)): + finished_request = finished_request or _coerce_finished_flag(request.is_finished()) + if not isinstance(pooling_output, dict): + if finished_request: + return { + "code_predictor_codes": [], + "finished": torch.tensor(True, dtype=torch.bool), + } + return None + + latent = pooling_output.get("latent_audio_feat") + if isinstance(latent, torch.Tensor) and latent.numel() == 0: + latent = None + + if latent is None: + if finished_request: + return { + "code_predictor_codes": [], + "finished": torch.tensor(True, dtype=torch.bool), + } + return None + + serialized_codes = _serialize_latent_to_codes(latent) + out: dict[str, Any] = { + "code_predictor_codes": serialized_codes, + "finished": torch.tensor(finished_request, dtype=torch.bool), + } + return out diff --git a/vllm_omni/platforms/npu/stage_configs/voxcpm.yaml b/vllm_omni/platforms/npu/stage_configs/voxcpm.yaml new file mode 100644 index 000000000..dcd1f4051 --- /dev/null +++ b/vllm_omni/platforms/npu/stage_configs/voxcpm.yaml @@ -0,0 +1,67 @@ +stage_args: + - stage_id: 0 + stage_type: llm + is_comprehension: true + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + dtype: bfloat16 + model_stage: latent_generator + model_arch: VoxCPMForConditionalGeneration + # Optional persistent HF-compatible config dir for native VoxCPM models. + hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: latent + gpu_memory_utilization: 0.75 + distributed_executor_backend: "mp" + max_num_batched_tokens: 4096 + max_model_len: 4096 + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 4096 + seed: 42 + detokenize: false + repetition_penalty: 1.0 + final_output: false + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + dtype: float32 + model_stage: vae + model_arch: VoxCPMForConditionalGeneration + hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio + gpu_memory_utilization: 0.1 + distributed_executor_backend: "mp" + max_num_batched_tokens: 8192 + max_model_len: 4096 + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.voxcpm.latent2vae + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 1 + seed: 42 + detokenize: true + repetition_penalty: 1.0 diff --git a/vllm_omni/platforms/npu/stage_configs/voxcpm_async_chunk.yaml b/vllm_omni/platforms/npu/stage_configs/voxcpm_async_chunk.yaml new file mode 100644 index 000000000..0a4ed7497 --- /dev/null +++ b/vllm_omni/platforms/npu/stage_configs/voxcpm_async_chunk.yaml @@ -0,0 +1,93 @@ +async_chunk: true +stage_args: + - stage_id: 0 + stage_type: llm + is_comprehension: true + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + dtype: bfloat16 + model_stage: latent_generator + model_arch: VoxCPMForConditionalGeneration + # Optional persistent HF-compatible config dir for native VoxCPM models. + hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: latent + gpu_memory_utilization: 0.75 + distributed_executor_backend: "mp" + max_num_batched_tokens: 4096 + max_model_len: 4096 + custom_process_next_stage_input_func: vllm_omni.model_executor.stage_input_processors.voxcpm.latent2vae_async_chunk + output_connectors: + to_stage_1: connector_of_shared_memory + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 4096 + seed: 42 + detokenize: false + repetition_penalty: 1.0 + final_output: false + + - stage_id: 1 + stage_type: llm + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + dtype: float32 + model_stage: vae + model_arch: VoxCPMForConditionalGeneration + hf_config_path: ${oc.env:VLLM_OMNI_VOXCPM_HF_CONFIG_PATH,} + worker_type: generation + scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler + enforce_eager: true + trust_remote_code: true + async_scheduling: false + enable_prefix_caching: false + engine_output_type: audio + gpu_memory_utilization: 0.1 + distributed_executor_backend: "mp" + max_num_batched_tokens: 8192 + max_model_len: 4096 + engine_input_source: [0] + input_connectors: + from_stage_0: connector_of_shared_memory + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 1 + seed: 42 + detokenize: true + repetition_penalty: 1.0 + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 + + connectors: + connector_of_shared_memory: + name: SharedMemoryConnector + extra: + shm_threshold_bytes: 65536 + codec_streaming: false + connector_get_sleep_s: 0.01 + connector_get_max_wait_first_chunk: 3000 + connector_get_max_wait: 300 + + edges: + - from: 0 + to: 1 + window_size: -1 diff --git a/vllm_omni/transformers_utils/configs/__init__.py b/vllm_omni/transformers_utils/configs/__init__.py index 5f957c2f6..0aa3624f8 100644 --- a/vllm_omni/transformers_utils/configs/__init__.py +++ b/vllm_omni/transformers_utils/configs/__init__.py @@ -17,6 +17,7 @@ _CLASS_TO_MODULE: dict[str, str] = { "FishSpeechConfig": "vllm_omni.transformers_utils.configs.fish_speech", "FishSpeechSlowARConfig": "vllm_omni.transformers_utils.configs.fish_speech", "FishSpeechFastARConfig": "vllm_omni.transformers_utils.configs.fish_speech", + "VoxCPMConfig": "vllm_omni.transformers_utils.configs.voxcpm", "VoxCPM2Config": "vllm_omni.transformers_utils.configs.voxcpm2", } @@ -28,6 +29,7 @@ __all__ = [ "FishSpeechConfig", "FishSpeechSlowARConfig", "FishSpeechFastARConfig", + "VoxCPMConfig", "VoxCPM2Config", ] @@ -49,4 +51,5 @@ def __dir__(): # run as soon as `vllm_omni.transformers_utils.configs` is imported. from vllm_omni.transformers_utils.configs import fish_speech as _fish_speech # noqa: F401, E402 from vllm_omni.transformers_utils.configs import mammoth_moda2 as _mammoth_moda2 # noqa: F401, E402 +from vllm_omni.transformers_utils.configs import voxcpm as _voxcpm # noqa: F401, E402 from vllm_omni.transformers_utils.configs import voxcpm2 as _voxcpm2 # noqa: F401, E402 diff --git a/vllm_omni/transformers_utils/configs/voxcpm.py b/vllm_omni/transformers_utils/configs/voxcpm.py new file mode 100644 index 000000000..026783891 --- /dev/null +++ b/vllm_omni/transformers_utils/configs/voxcpm.py @@ -0,0 +1,68 @@ +from transformers import AutoConfig +from transformers.configuration_utils import PretrainedConfig + + +class VoxCPMConfig(PretrainedConfig): + model_type = "voxcpm" + keys_to_ignore_at_inference = ["past_key_values"] + + def __init__( + self, + bos_token_id: int = 1, + eos_token_id: int = 2, + vocab_size: int = 32000, + hidden_size: int = 1024, + intermediate_size: int = 4096, + max_position_embeddings: int = 4096, + num_attention_heads: int = 16, + num_hidden_layers: int = 24, + num_key_value_heads: int = 16, + rms_norm_eps: float = 1e-6, + rope_theta: float = 10000.0, + rope_scaling: dict | None = None, + lm_config: dict | None = None, + encoder_config: dict | None = None, + dit_config: dict | None = None, + audio_vae_config: dict | None = None, + patch_size: int = 2, + feat_dim: int = 64, + residual_lm_num_layers: int = 6, + scalar_quantization_latent_dim: int = 256, + scalar_quantization_scale: int = 9, + max_length: int = 4096, + device: str = "cuda", + dtype: str = "bfloat16", + dit_mean_mode: bool = False, + **kwargs, + ): + super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs) + self.vocab_size = vocab_size + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.max_position_embeddings = max_position_embeddings + self.num_attention_heads = num_attention_heads + self.num_hidden_layers = num_hidden_layers + self.num_key_value_heads = num_key_value_heads + self.rms_norm_eps = rms_norm_eps + self.rope_theta = rope_theta + self.rope_scaling = rope_scaling + + self.lm_config = lm_config or {} + self.encoder_config = encoder_config or {} + self.dit_config = dit_config or {} + self.audio_vae_config = audio_vae_config + + self.patch_size = patch_size + self.feat_dim = feat_dim + self.residual_lm_num_layers = residual_lm_num_layers + self.scalar_quantization_latent_dim = scalar_quantization_latent_dim + self.scalar_quantization_scale = scalar_quantization_scale + self.max_length = max_length + self.device = device + self.dtype = dtype + self.dit_mean_mode = dit_mean_mode + + +AutoConfig.register("voxcpm", VoxCPMConfig) + +__all__ = ["VoxCPMConfig"] From 82f8c93343552d81e0e4730d90ce08e072fc3bcb Mon Sep 17 00:00:00 2001 From: Juan Pablo Zuluaga <46724788+JuanPZuluaga@users.noreply.github.com> Date: Wed, 15 Apr 2026 09:14:57 +0200 Subject: [PATCH 13/97] [Feat][Qwen3-Omni] Shared code predictor module for Qwen3-TTS and Qwen3-Omni (#2375) Signed-off-by: JuanPZuluaga Co-authored-by: Hongsheng Liu --- .../qwen3_tts/test_code_predictor_dtype.py | 92 ++- vllm_omni/engine/stage_init_utils.py | 5 +- .../model_executor/models/common/__init__.py | 0 .../models/common/qwen3_code_predictor.py | 654 ++++++++++++++++++ .../qwen3_omni_moe_code_predictor_mtp.py | 520 +------------- .../qwen3_tts_code_predictor_vllm.py | 571 +-------------- 6 files changed, 778 insertions(+), 1064 deletions(-) create mode 100644 vllm_omni/model_executor/models/common/__init__.py create mode 100644 vllm_omni/model_executor/models/common/qwen3_code_predictor.py diff --git a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py index b0ce10a8d..8798cb3ca 100644 --- a/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py +++ b/tests/model_executor/models/qwen3_tts/test_code_predictor_dtype.py @@ -21,7 +21,7 @@ import torch from pytest_mock import MockerFixture # Direct file import to avoid vllm_omni.__init__ patch dependencies. -_BASE = os.path.join( +_MODELS = os.path.join( os.path.dirname(__file__), os.pardir, os.pardir, @@ -30,14 +30,16 @@ _BASE = os.path.join( "vllm_omni", "model_executor", "models", - "qwen3_tts", ) +_BASE = os.path.join(_MODELS, "qwen3_tts") +_COMMON = os.path.join(_MODELS, "common") def _load_module(name: str, filename: str): path = os.path.abspath(os.path.join(_BASE, filename)) spec = importlib.util.spec_from_file_location(name, path) mod = importlib.util.module_from_spec(spec) + sys.modules[name] = mod # register before exec (needed for dataclasses etc.) spec.loader.exec_module(mod) return mod @@ -59,8 +61,17 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]: weight_utils_mock = mocker.MagicMock() weight_utils_mock.default_weight_loader = lambda p, w: None - pkg = types.ModuleType("vllm_omni.model_executor.models.qwen3_tts") - pkg.__path__ = [os.path.abspath(_BASE)] + tts_pkg = types.ModuleType("vllm_omni.model_executor.models.qwen3_tts") + tts_pkg.__path__ = [os.path.abspath(_BASE)] + + common_pkg = types.ModuleType("vllm_omni.model_executor.models.common") + common_pkg.__path__ = [os.path.abspath(_COMMON)] + + models_pkg = types.ModuleType("vllm_omni.model_executor.models") + models_pkg.__path__ = [os.path.abspath(_MODELS)] + + vllm_parallel_mock = mocker.MagicMock() + vllm_parallel_mock.VocabParallelEmbedding = torch.nn.Embedding return { "vllm_omni": mocker.MagicMock(), @@ -69,9 +80,11 @@ def _build_mock_modules(mocker: MockerFixture) -> dict[str, object]: "vllm.config": mocker.MagicMock(), "vllm.config.vllm": vllm_config_mod, "vllm.model_executor.model_loader.weight_utils": weight_utils_mock, + "vllm.model_executor.layers.vocab_parallel_embedding": vllm_parallel_mock, "vllm_omni.model_executor": types.ModuleType("vllm_omni.model_executor"), - "vllm_omni.model_executor.models": types.ModuleType("vllm_omni.model_executor.models"), - "vllm_omni.model_executor.models.qwen3_tts": pkg, + "vllm_omni.model_executor.models": models_pkg, + "vllm_omni.model_executor.models.common": common_pkg, + "vllm_omni.model_executor.models.qwen3_tts": tts_pkg, } @@ -88,6 +101,15 @@ def _load_target_classes(mocker: MockerFixture): ) sys.modules["vllm_omni.model_executor.models.qwen3_tts.configuration_qwen3_tts"] = config_mod + # Load the shared common module (thin wrappers import from it) + common_cp_path = os.path.abspath(os.path.join(_COMMON, "qwen3_code_predictor.py")) + common_spec = importlib.util.spec_from_file_location( + "vllm_omni.model_executor.models.common.qwen3_code_predictor", common_cp_path + ) + common_cp_mod = importlib.util.module_from_spec(common_spec) + sys.modules["vllm_omni.model_executor.models.common.qwen3_code_predictor"] = common_cp_mod + common_spec.loader.exec_module(common_cp_mod) + cp_mod = _load_module( "vllm_omni.model_executor.models.qwen3_tts.qwen3_tts_code_predictor_vllm", "qwen3_tts_code_predictor_vllm.py", @@ -104,6 +126,7 @@ def loaded_target_classes(mocker: MockerFixture): config_mod.Qwen3TTSTalkerConfig, cp_mod.Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM, cp_mod.Qwen3TTSTalkerCodePredictorModelVLLM, + cp_mod.CodePredictorWrapperConfig, ) @@ -114,6 +137,7 @@ def _make_tiny_config(loaded_target_classes) -> tuple: qwen3_tts_talker_config, _, _, + _, ) = loaded_target_classes cp_config = qwen3_tts_talker_code_predictor_config( vocab_size=64, @@ -145,7 +169,7 @@ class TestCodePredictorDtypeAlignment: def test_ensure_buffers_uses_given_dtype(self, mocker: MockerFixture, loaded_target_classes) -> None: """_ensure_buffers should create proj_buf with the given dtype.""" - _, _, code_predictor_wrapper, _ = loaded_target_classes + _, _, code_predictor_wrapper, _, _ = loaded_target_classes cp_config, talker_config = _make_tiny_config(loaded_target_classes) vllm_config = _make_vllm_config(mocker) @@ -156,17 +180,17 @@ class TestCodePredictorDtypeAlignment: ) # Create buffer in float16 - predictor._ensure_buffers(torch.device("cpu"), torch.float16) + predictor._ensure_buffers(torch.device("cpu"), torch.float16, 4) assert predictor._proj_buf is not None assert predictor._proj_buf.dtype == torch.float16 # Re-create buffer in float32 (different dtype triggers re-allocation) - predictor._ensure_buffers(torch.device("cpu"), torch.float32) + predictor._ensure_buffers(torch.device("cpu"), torch.float32, 4) assert predictor._proj_buf.dtype == torch.float32 def test_warmup_aligns_buffer_to_model_params(self, mocker: MockerFixture, loaded_target_classes) -> None: """_warmup_buckets should align proj_buf dtype to model parameters.""" - _, _, code_predictor_wrapper, _ = loaded_target_classes + _, _, code_predictor_wrapper, _, _ = loaded_target_classes cp_config, talker_config = _make_tiny_config(loaded_target_classes) vllm_config = _make_vllm_config(mocker, max_num_seqs=2) @@ -180,7 +204,7 @@ class TestCodePredictorDtypeAlignment: predictor = predictor.to(torch.float16) # Pre-create proj_buf with WRONG dtype (float32) — simulating the bug - predictor._ensure_buffers(torch.device("cpu"), torch.float32) + predictor._ensure_buffers(torch.device("cpu"), torch.float32, 2) assert predictor._proj_buf.dtype == torch.float32 # Simulate _setup_compile having cached model dtype and compiled forward @@ -194,7 +218,7 @@ class TestCodePredictorDtypeAlignment: def test_setup_compile_caches_model_dtype(self, mocker: MockerFixture, loaded_target_classes) -> None: """_setup_compile should cache model parameter dtype.""" - _, _, code_predictor_wrapper, _ = loaded_target_classes + _, _, code_predictor_wrapper, _, _ = loaded_target_classes cp_config, talker_config = _make_tiny_config(loaded_target_classes) vllm_config = _make_vllm_config(mocker, max_num_seqs=2) @@ -211,7 +235,7 @@ class TestCodePredictorDtypeAlignment: def test_forward_with_mismatched_input_dtype(self, mocker: MockerFixture, loaded_target_classes) -> None: """forward() should not crash when inputs are float32 but model is float16.""" - _, _, code_predictor_wrapper, _ = loaded_target_classes + _, _, code_predictor_wrapper, _, _ = loaded_target_classes cp_config, talker_config = _make_tiny_config(loaded_target_classes) vllm_config = _make_vllm_config(mocker, max_num_seqs=2) @@ -250,9 +274,9 @@ class TestCodePredictorModelDtype: def test_model_forward_float16(self, loaded_target_classes) -> None: """Inner model forward should work in float16.""" - _, _, _, code_predictor_model = loaded_target_classes + _, _, _, code_predictor_model, _ = loaded_target_classes cp_config, _ = _make_tiny_config(loaded_target_classes) - model = code_predictor_model(cp_config, talker_hidden_size=32).to(torch.float16) + model = code_predictor_model(cp_config, embedding_dim=32).to(torch.float16) bsz, seq_len = 1, 4 inputs = torch.randn(bsz, seq_len, 32, dtype=torch.float16) @@ -264,9 +288,9 @@ class TestCodePredictorModelDtype: def test_model_forward_float32(self, loaded_target_classes) -> None: """Inner model forward should work in float32.""" - _, _, _, code_predictor_model = loaded_target_classes + _, _, _, code_predictor_model, _ = loaded_target_classes cp_config, _ = _make_tiny_config(loaded_target_classes) - model = code_predictor_model(cp_config, talker_hidden_size=32).to(torch.float32) + model = code_predictor_model(cp_config, embedding_dim=32).to(torch.float32) bsz, seq_len = 1, 4 inputs = torch.randn(bsz, seq_len, 32, dtype=torch.float32) @@ -275,3 +299,37 @@ class TestCodePredictorModelDtype: output = model(inputs, pos_ids) assert output.dtype == torch.float32 assert output.shape == (bsz, seq_len, 32) + + +class TestCodePredictorWrapperConfig: + """Test wrapper configuration for different models.""" + + def test_omni_config(self, loaded_target_classes) -> None: + """Qwen3-Omni uses correct wrapper config.""" + _, _, _, _, code_predictor_wrapper_config = loaded_target_classes + config = code_predictor_wrapper_config( + use_cuda_graphs=False, + use_parallel_embedding=True, + use_projection=False, + return_proj_buf=True, + sampling_mode="stored", + ) + assert config.use_cuda_graphs is False + assert config.use_parallel_embedding is True + assert config.return_proj_buf is True + assert config.sampling_mode == "stored" + + def test_tts_config(self, loaded_target_classes) -> None: + """Qwen3-TTS uses correct wrapper config.""" + _, _, _, _, code_predictor_wrapper_config = loaded_target_classes + config = code_predictor_wrapper_config( + use_cuda_graphs=True, + use_parallel_embedding=False, + use_projection=True, + return_proj_buf=False, + sampling_mode="per_call", + ) + assert config.use_cuda_graphs is True + assert config.use_parallel_embedding is False + assert config.return_proj_buf is False + assert config.sampling_mode == "per_call" diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index bf40aa77c..3a7fe4bad 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -192,8 +192,9 @@ def extract_stage_metadata(stage_config: Any) -> StageMetadata: default_sampling_params: OmniSamplingParams = SPClass(**default_sp) custom_process_input_func: Callable | None = None - if hasattr(stage_config, "custom_process_input_func"): - mod_path, fn_name = stage_config.custom_process_input_func.rsplit(".", 1) + _cpif_path = getattr(stage_config, "custom_process_input_func", None) + if _cpif_path: + mod_path, fn_name = _cpif_path.rsplit(".", 1) custom_process_input_func = getattr(importlib.import_module(mod_path), fn_name) prompt_expand_func: Callable | None = None diff --git a/vllm_omni/model_executor/models/common/__init__.py b/vllm_omni/model_executor/models/common/__init__.py new file mode 100644 index 000000000..e69de29bb diff --git a/vllm_omni/model_executor/models/common/qwen3_code_predictor.py b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py new file mode 100644 index 000000000..3a904442f --- /dev/null +++ b/vllm_omni/model_executor/models/common/qwen3_code_predictor.py @@ -0,0 +1,654 @@ +"""Qwen3 Code Predictor -- optimized re-prefill, no KV cache. + +Shared by Qwen3-Omni and Qwen3-TTS talker models. + +* SDPA attention (F.scaled_dot_product_attention) with native GQA support +* HF-compatible numerics (float32 RMSNorm, float32 RoPE, separate linear layers) +* Per-call embedding buffer to avoid cross-request aliasing +* Pre-allocated position_ids (read-only, safe to persist) +* torch.compile (epilogue_fusion=False) on inner transformer by default +* Optional manual CUDA graph capture per batch-size bucket +* Inline sampling (top-k + top-p) -- no custom op overhead +""" + +from __future__ import annotations + +import dataclasses +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.platforms import current_omni_platform + +logger = init_logger(__name__) + + +# =================================================================== +# HF-numerics-compatible layers for code predictor +# =================================================================== +# +# These use plain PyTorch ops (nn.Linear, manual RMSNorm in float32, +# rotate_half RoPE) to produce outputs numerically identical to the +# HuggingFace reference. vLLM's fused kernels (RMSNorm, QKVParallel, +# get_rope) introduce small precision differences that compound across +# the autoregressive steps of the code predictor, causing severe +# audio quality degradation. +# +# See: https://github.com/vllm-project/vllm-omni/issues/2274 + + +class _RMSNorm(nn.Module): + """RMSNorm matching HuggingFace's implementation exactly. + + Computes variance in float32 to avoid bfloat16 precision loss. + """ + + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return self.weight * hidden_states.to(input_dtype) + + +def _rotate_half(x: torch.Tensor) -> torch.Tensor: + """Rotates half the hidden dims of the input.""" + x1 = x[..., : x.shape[-1] // 2] + x2 = x[..., x.shape[-1] // 2 :] + return torch.cat((-x2, x1), dim=-1) + + +class _RotaryEmbedding(nn.Module): + """RoPE matching HuggingFace's implementation exactly. + + Forces float32 computation for cos/sin, matching HF's torch.autocast(enabled=False). + """ + + def __init__(self, config) -> None: + super().__init__() + head_dim = getattr( + config, + "head_dim", + config.hidden_size // config.num_attention_heads, + ) + rope_theta = getattr(config, "rope_theta", 10000.0) + inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + # position_ids: [batch, seq_len] + inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position_ids_expanded = position_ids[:, None, :].float() + + # Force float32 (matching HF) + device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" + with torch.autocast(device_type=device_type, enabled=False): + freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos() + sin = emb.sin() + + return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) + + +# =================================================================== +# Attention +# =================================================================== + + +class CodePredictorAttention(nn.Module): + """Multi-head self-attention for code predictor. + + Uses ``F.scaled_dot_product_attention`` with HF-compatible RoPE and RMSNorm. + No KV cache -- the code predictor always re-prefills the full (short) + sequence each AR step. + + Input : [B, seq_len, hidden_size] + Output: [B, seq_len, hidden_size] + """ + + def __init__(self, config, *, prefix: str = "") -> None: + super().__init__() + self.num_heads = config.num_attention_heads + self.num_kv_heads = config.num_key_value_heads + self.head_dim = getattr( + config, + "head_dim", + config.hidden_size // config.num_attention_heads, + ) + self.hidden_size = config.hidden_size + self.scaling = self.head_dim**-0.5 + self._use_gqa = self.num_kv_heads != self.num_heads + + # Separate q/k/v projections matching HF (no fused packing) + bias = getattr(config, "attention_bias", False) + self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=bias) + self.k_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=bias) + self.v_proj = nn.Linear(self.hidden_size, self.num_kv_heads * self.head_dim, bias=bias) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + bsz, seq_len, _ = hidden_states.shape + hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim) + hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim) + + q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2) + k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2) + v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2) + + cos, sin = position_embeddings + # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads + cos = cos.unsqueeze(1) # [batch, 1, seq_len, head_dim] + sin = sin.unsqueeze(1) + q = (q * cos) + (_rotate_half(q) * sin) + k = (k * cos) + (_rotate_half(k) * sin) + + attn_out = F.scaled_dot_product_attention( + q, + k, + v, + scale=self.scaling, + is_causal=True, + enable_gqa=self._use_gqa, + ) + + attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) + return self.o_proj(attn_out) + + +# =================================================================== +# MLP +# =================================================================== + + +class CodePredictorMLP(nn.Module): + """SiLU-gated MLP for code predictor, matching HF's implementation.""" + + def __init__(self, config, *, prefix: str = "") -> None: + super().__init__() + self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) + self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) + + +# =================================================================== +# Decoder Layer +# =================================================================== + + +class CodePredictorDecoderLayer(nn.Module): + """Transformer decoder layer (SDPA, no KV cache).""" + + def __init__(self, config, *, prefix: str = "") -> None: + super().__init__() + self.self_attn = CodePredictorAttention(config, prefix=f"{prefix}.self_attn") + self.mlp = CodePredictorMLP(config, prefix=f"{prefix}.mlp") + self.input_layernorm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + ) -> torch.Tensor: + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn(hidden_states, position_embeddings) + hidden_states = residual + hidden_states + + residual = hidden_states + hidden_states = self.post_attention_layernorm(hidden_states) + hidden_states = self.mlp(hidden_states) + hidden_states = residual + hidden_states + return hidden_states + + +# =================================================================== +# Base Transformer Model (re-prefill, no KV cache) +# =================================================================== + + +class CodePredictorBaseModel(nn.Module): + """Inner transformer for code predictor. + + Signature: ``forward(inputs_embeds, position_ids) -> hidden_states`` + """ + + def __init__( + self, + config, + *, + embedding_dim: int | None = None, + use_parallel_embedding: bool = False, + prefix: str = "", + ) -> None: + super().__init__() + self.config = config + + emb_dim = int(embedding_dim) if embedding_dim is not None else int(config.hidden_size) + if use_parallel_embedding: + self.codec_embedding = nn.ModuleList( + [VocabParallelEmbedding(config.vocab_size, emb_dim) for _ in range(config.num_code_groups - 1)] + ) + else: + self.codec_embedding = nn.ModuleList( + [nn.Embedding(config.vocab_size, emb_dim) for _ in range(config.num_code_groups - 1)] + ) + + self.layers = nn.ModuleList( + [ + CodePredictorDecoderLayer(config, prefix=f"{prefix}.layers.{idx}") + for idx in range(config.num_hidden_layers) + ] + ) + self.norm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = _RotaryEmbedding(config) + + def get_input_embeddings(self) -> nn.ModuleList: + return self.codec_embedding + + def forward( + self, + inputs_embeds: torch.Tensor, + position_ids: torch.Tensor, + ) -> torch.Tensor: + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + for layer in self.layers: + hidden_states = layer(hidden_states, position_embeddings) + hidden_states = self.norm(hidden_states) + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + for name, loaded_weight in weights: + if "rotary_emb.inv_freq" in name: + continue + param = params_dict.get(name) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + return loaded_params + + +# =================================================================== +# Wrapper Configuration +# =================================================================== + + +@dataclasses.dataclass +class CodePredictorWrapperConfig: + """Controls behavioral differences between model-specific code predictors.""" + + use_cuda_graphs: bool = False + use_parallel_embedding: bool = False + use_projection: bool = False + return_proj_buf: bool = False + sampling_mode: str = "stored" + + +# =================================================================== +# Code Predictor Wrapper (optimized re-prefill, persistent buffers) +# =================================================================== + + +class CodePredictorWrapper(nn.Module): + """Optimized code predictor -- re-prefill approach, no KV cache. + + Each AR step forwards the full growing sequence (len 2 -> num_code_groups+1) + through the transformer. The extra O(T^2) FLOPs are negligible for + short sequences, and this avoids all KV-cache management overhead. + + Optimizations: + 1. Per-call embedding buffer -- avoids cross-request aliasing. + 2. Pre-allocated position_ids -- no torch.arange per step. + 3. Cached module references -- bypass ModuleList indexing. + 4. torch.compile on inner transformer. + 5. Inline sampling (top-k + top-p) -- no custom op overhead. + 6. Optional manual CUDA graph capture per batch-size bucket. + """ + + def __init__( + self, + *, + vllm_config: VllmConfig, + cp_config, + wrapper_config: CodePredictorWrapperConfig, + talker_hidden_size: int | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self._vllm_config = vllm_config + self.config = cp_config + self._wrapper_config = wrapper_config + self.prefix = prefix + + self._num_groups = int(cp_config.num_code_groups) + self._cp_hidden = int(cp_config.hidden_size) + + # For Omni backward compat (accessed by the talker) + self.num_code_groups = self._num_groups + + # Determine embedding dimension + _talker_hidden = int(talker_hidden_size) if talker_hidden_size is not None else self._cp_hidden + + self.model = CodePredictorBaseModel( + cp_config, + embedding_dim=_talker_hidden, + use_parallel_embedding=wrapper_config.use_parallel_embedding, + prefix=f"{prefix}.model" if prefix else "model", + ) + + self.lm_head = nn.ModuleList( + [nn.Linear(cp_config.hidden_size, cp_config.vocab_size, bias=False) for _ in range(self._num_groups - 1)] + ) + + # Projection: Identity when hidden sizes match or not needed + if wrapper_config.use_projection and _talker_hidden != self._cp_hidden: + self.small_to_mtp_projection = nn.Linear(_talker_hidden, self._cp_hidden, bias=True) + else: + self.small_to_mtp_projection = nn.Identity() + + # Sampling defaults for "stored" mode + self._top_k: int = 50 + self._top_p: float = 0.8 + + # Lazily initialised state + self._proj_buf: torch.Tensor | None = None + self._model_dtype: torch.dtype | None = None + self._compiled_model_fwd = None + self._bucket_sizes: list[int] = [] + self._bucket_pos_ids: dict[int, torch.Tensor] = {} + self._lm_heads_list: list[nn.Module] | None = None + self._codec_embeds_list: list[nn.Module] | None = None + self._cuda_graphs: dict[int, tuple[torch.cuda.CUDAGraph, torch.Tensor]] = {} + + def get_input_embeddings(self) -> nn.ModuleList: + return self.model.get_input_embeddings() + + def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8) -> None: + """Configure sampling parameters to maintain consistency with previous implementation.""" + self._top_k = top_k + self._top_p = top_p + logger.debug("Sampling parameters updated: top_k=%d, top_p=%.2f", top_k, top_p) + + # ------------------------------------------------------------------ + # Lazy-init helpers + # ------------------------------------------------------------------ + + def _ensure_buffers(self, device: torch.device, dtype: torch.dtype, bsz: int) -> None: + """Ensure the projection buffer can hold at least *bsz* rows.""" + max_seq = self._num_groups + 1 + if ( + self._proj_buf is not None + and self._proj_buf.device == device + and self._proj_buf.dtype == dtype + and self._proj_buf.shape[0] >= bsz + ): + return + self._proj_buf = torch.zeros(bsz, max_seq, self._cp_hidden, dtype=dtype, device=device) + + def _setup_compile(self) -> None: + """Lazily set up torch.compile with optional CUDA graph capture.""" + if self._compiled_model_fwd is not None: + return + + # Cache model parameter dtype so forward() doesn't need to query it + # on every call. Also ensures warmup buffers match model precision + # even when upstream modules produce a different dtype (#2385). + self._model_dtype = next(self.model.parameters()).dtype + self._lm_heads_list = list(self.lm_head) + self._codec_embeds_list = list(self.model.codec_embedding) + + if not current_omni_platform.supports_torch_inductor(): + logger.warning_once("code_predictor: torch.compile disabled") + self._compiled_model_fwd = self.model.forward + return + + # torch.compile fuses RMSNorm/RoPE in ways that lose float32 + # precision, compounding across AR steps. Use epilogue_fusion=False + # to disable the problematic fusions while still getting kernel + # fusion benefits for the linear layers and SDPA. + self._compiled_model_fwd = torch.compile( + self.model.forward, + dynamic=False, + options={"epilogue_fusion": False}, + ) + self._warmup_buckets() + + if self._wrapper_config.use_cuda_graphs: + self._capture_cuda_graphs() + logger.info("code_predictor: torch.compile (no epilogue fusion) + CUDA graphs") + else: + logger.info("code_predictor: torch.compile (dynamic=False, no epilogue fusion)") + + def _padded_bsz(self, bsz: int) -> int: + """Round batch size up to nearest power-of-2 bucket.""" + for bucket in self._bucket_sizes: + if bsz <= bucket: + return bucket + return bsz + + def _warmup_buckets(self) -> None: + """Warmup power-of-2 batch-size buckets to front-load Inductor compilation.""" + max_bsz = self._vllm_config.scheduler_config.max_num_seqs + bucket_sizes = [1 << i for i in range(max_bsz.bit_length()) if (1 << i) <= max_bsz] + if max_bsz not in bucket_sizes: + bucket_sizes.append(max_bsz) + self._bucket_sizes = sorted(bucket_sizes) + + max_seq = self._num_groups + 1 + device = next(self.model.parameters()).device + + # Ensure proj_buf matches model parameter dtype to avoid dtype + # mismatch during warmup compilation (see #2385). + self._ensure_buffers(device, self._model_dtype, max(self._bucket_sizes)) + proj_buf = self._proj_buf + + for bsz in self._bucket_sizes: + pos_ids = torch.arange(max_seq, device=device, dtype=torch.long).unsqueeze(0).expand(bsz, -1).contiguous() + self._bucket_pos_ids[bsz] = pos_ids + for _ in range(3): + self._compiled_model_fwd(proj_buf[:bsz, :max_seq, :], pos_ids) + logger.info("code_predictor: warmup done for buckets %s", self._bucket_sizes) + + def _capture_cuda_graphs(self) -> None: + """Capture a CUDA graph per bucket using vLLM's global graph pool.""" + from vllm.platforms import current_platform + + pool = current_platform.get_global_graph_pool() + max_seq = self._num_groups + 1 + proj_buf = self._proj_buf + + for bsz in self._bucket_sizes: + static_input = proj_buf[:bsz, :max_seq, :] + pos_ids = self._bucket_pos_ids[bsz] + + g = torch.cuda.CUDAGraph() + with torch.cuda.graph(g, pool=pool): + static_output = self._compiled_model_fwd(static_input, pos_ids) + + self._cuda_graphs[bsz] = (g, static_output) + + logger.info("code_predictor: captured CUDA graphs for buckets %s", self._bucket_sizes) + + # ------------------------------------------------------------------ + # Forward -- re-prefill + inline sampling + # ------------------------------------------------------------------ + + @torch.inference_mode() + def forward( + self, + layer0_code: torch.Tensor, + layer0_embed: torch.Tensor, + last_talker_hidden: torch.Tensor, + do_sample: bool = True, + temperature: float = 0.9, + top_k: int = 50, + top_p: float = 1.0, + ) -> torch.Tensor | tuple[torch.Tensor, torch.Tensor]: + """Predict residual codebooks 1..G-1 autoregressively via re-prefill.""" + bsz = int(layer0_code.shape[0]) + num_groups = self._num_groups + device = layer0_code.device + + # _setup_compile caches _model_dtype on first call; use it for buffers + # so they always match model weight precision (#2385). + self._setup_compile() + dtype = self._model_dtype + + padded_bsz = self._padded_bsz(bsz) + self._ensure_buffers(device, dtype, padded_bsz) + + proj_buf = self._proj_buf + max_seq = num_groups + 1 + projection = self.small_to_mtp_projection + model_fwd = self._compiled_model_fwd + lm_heads = self._lm_heads_list + codec_embeds = self._codec_embeds_list + + # Zero the padded region of the buffer + proj_buf[:padded_bsz].zero_() + + # Fill buffer positions 0 (talker hidden) & 1 (layer0 embed) + proj_buf[:bsz, 0, :] = projection(last_talker_hidden.reshape(bsz, 1, -1).to(dtype)).reshape(bsz, -1) + proj_buf[:bsz, 1, :] = projection(layer0_embed.reshape(bsz, 1, -1).to(dtype)).reshape(bsz, -1) + + # Get pre-computed pos_ids for this bucket + full_pos_ids = self._bucket_pos_ids.get(padded_bsz) + if full_pos_ids is None: + full_pos_ids = ( + torch.arange(max_seq, device=device, dtype=torch.long).unsqueeze(0).expand(padded_bsz, -1).contiguous() + ) + + # Use captured CUDA graph if available, otherwise call compiled fn. + cuda_graph_entry = self._cuda_graphs.get(padded_bsz) + + # Prepare sampling parameters + stored_mode = self._wrapper_config.sampling_mode == "stored" + if stored_mode: + s_top_k = self._top_k + s_top_p = self._top_p + else: + use_sampling = do_sample and temperature > 0 + inv_temperature = 1.0 / max(temperature, 1e-6) if use_sampling else 0.0 + if use_sampling and top_p != 1.0: + raise NotImplementedError( + "top_p sampling is not implemented for the vLLM-native code predictor; please set top_p=1.0." + ) + + # Output codes -- shape depends on return mode + if self._wrapper_config.return_proj_buf: + all_codes = torch.empty(bsz, num_groups, 1, dtype=torch.int64, device=device) + all_codes[:, 0] = layer0_code.reshape(bsz, -1)[:, :1] + else: + all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device) + all_codes[:, 0] = layer0_code.reshape(bsz) + + # Autoregressive loop: predict layers 1..G-1 + for step in range(1, num_groups): + # Run transformer (CUDA graph replay or compiled forward) + if cuda_graph_entry is not None: + cuda_graph_entry[0].replay() + hidden_out = cuda_graph_entry[1] + else: + hidden_out = model_fwd(proj_buf[:padded_bsz, :max_seq, :], full_pos_ids) + + logits = lm_heads[step - 1](hidden_out[:bsz, step, :]) + + # Sample next code + if stored_mode: + # "stored" mode: top-k -> top-p -> softmax -> multinomial + if s_top_k > 0: + topk_vals, _ = logits.topk(s_top_k, dim=-1) + logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf")) + if s_top_p < 1.0: + sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True) + sorted_probs = F.softmax(sorted_logits, dim=-1) + cumulative_probs = sorted_probs.cumsum(dim=-1) + remove_mask = (cumulative_probs - sorted_probs) >= s_top_p + sorted_logits[remove_mask] = float("-inf") + logits = sorted_logits.scatter(1, sorted_idx, sorted_logits) + probs = F.softmax(logits, dim=-1) + code = torch.multinomial(probs, num_samples=1) + else: + # "per_call" mode: temperature-scaled + top-k + if use_sampling: + scaled = logits * inv_temperature + if top_k > 0: + topk_vals, _ = scaled.topk(top_k, dim=-1) + scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) + probs = F.softmax(scaled, dim=-1) + code = torch.multinomial(probs, num_samples=1) + else: + code = logits.argmax(dim=-1, keepdim=True) + + # Store code + if self._wrapper_config.return_proj_buf: + all_codes[:, step] = code + else: + all_codes[:, step] = code.reshape(bsz) + + # Embed predicted code -> project -> next buffer position + if step < num_groups - 1 or self._wrapper_config.return_proj_buf: + new_embed = codec_embeds[step - 1](code) + proj_buf[:bsz, step + 1, :] = projection(new_embed.reshape(bsz, 1, -1)).reshape(bsz, -1) + + if self._wrapper_config.return_proj_buf: + return all_codes, proj_buf[:bsz].clone() + return all_codes + + # ------------------------------------------------------------------ + # Weight loading + # ------------------------------------------------------------------ + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights directly (no fused projection remapping needed).""" + loaded: set[str] = set() + model_weights: list[tuple[str, torch.Tensor]] = [] + other_weights: list[tuple[str, torch.Tensor]] = [] + + for name, w in weights: + if "rotary_emb.inv_freq" in name: + continue + if name.startswith("model."): + model_weights.append((name[len("model.") :], w)) + else: + other_weights.append((name, w)) + + loaded_model = self.model.load_weights(model_weights) + loaded |= {f"model.{n}" for n in loaded_model} + + params = dict(self.named_parameters(remove_duplicate=False)) + for name, w in other_weights: + param = params.get(name) + if param is None: + continue + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, w) + loaded.add(name) + + return loaded diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py index 2ceaafdb6..819e22e18 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py @@ -1,510 +1,28 @@ -"""Qwen3-Omni Code Predictor -- optimized re-prefill, no KV cache. +"""Qwen3-Omni Code Predictor -- thin wrapper over CodePredictorWrapper.""" -* SDPA attention (F.scaled_dot_product_attention) with native GQA support -* HF-compatible numerics (float32 RMSNorm, float32 RoPE, separate linear layers) -* Per-call embedding buffer to avoid cross-request aliasing -* Pre-allocated position_ids (read-only, safe to persist) -* torch.compile (epilogue_fusion=False) on inner transformer by default -* Inline sampling (top-k + top-p) -- no custom op overhead -""" - -import torch -import torch.nn as nn -import torch.nn.functional as F from vllm.config import VllmConfig -from vllm.logger import init_logger -from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding -from vllm.model_executor.model_loader.weight_utils import default_weight_loader -from vllm_omni.platforms import current_omni_platform - -logger = init_logger(__name__) +from vllm_omni.model_executor.models.common.qwen3_code_predictor import ( + CodePredictorWrapper, + CodePredictorWrapperConfig, +) -# =================================================================== -# HF-numerics-compatible layers for code predictor -# =================================================================== -# -# These use plain PyTorch ops (nn.Linear, manual RMSNorm in float32, -# rotate_half RoPE) to produce outputs numerically identical to the -# HuggingFace reference. vLLM's fused kernels (RMSNorm, QKVParallel, -# get_rope) introduce small precision differences that compound across -# the autoregressive steps of the code predictor, causing severe -# audio quality degradation. -# -# See: https://github.com/vllm-project/vllm-omni/issues/2274 +class Qwen3OmniMoeTalkerCodePredictor(CodePredictorWrapper): + """Qwen3-Omni code predictor (no CUDA graphs, VocabParallelEmbedding).""" - -class _RMSNorm(nn.Module): - """RMSNorm matching HuggingFace's implementation exactly. - - Computes variance in float32 to avoid bfloat16 precision loss. - """ - - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) - - -def _rotate_half(x: torch.Tensor) -> torch.Tensor: - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -class _RotaryEmbedding(nn.Module): - """RoPE matching HuggingFace's implementation exactly. - - Forces float32 computation for cos/sin, matching HF's torch.autocast(enabled=False). - """ - - def __init__(self, config) -> None: - super().__init__() - head_dim = getattr( - config, - "head_dim", - config.hidden_size // config.num_attention_heads, - ) - rope_theta = getattr(config, "rope_theta", 10000.0) - inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - - def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - # position_ids: [batch, seq_len] - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) - position_ids_expanded = position_ids[:, None, :].float() - - # Force float32 (matching HF) - device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" - with torch.autocast(device_type=device_type, enabled=False): - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1) - cos = emb.cos() - sin = emb.sin() - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) - - -class Qwen3OmniCodePredictorAttention(nn.Module): - """Multi-head self-attention for code predictor. - - Uses ``F.scaled_dot_product_attention`` with HF-compatible RoPE and RMSNorm. - No KV cache -- the code predictor always re-prefills the full (short) - sequence each AR step. - - Input : [B, seq_len, hidden_size] - Output: [B, seq_len, hidden_size] - """ - - def __init__( - self, - config, - prefix: str = "", - ): - super().__init__() - cp_cfg = config.code_predictor_config - self.num_heads = cp_cfg.num_attention_heads - self.num_kv_heads = cp_cfg.num_key_value_heads - self.head_dim = getattr( - cp_cfg, - "head_dim", - cp_cfg.hidden_size // cp_cfg.num_attention_heads, - ) - self.hidden_size = cp_cfg.hidden_size - self.scaling = self.head_dim**-0.5 - self._use_gqa = self.num_kv_heads != self.num_heads - - # Separate q/k/v projections matching HF (no fused packing) - self.q_proj = nn.Linear( - self.hidden_size, - self.num_heads * self.head_dim, - bias=False, - ) - self.k_proj = nn.Linear( - self.hidden_size, - self.num_kv_heads * self.head_dim, - bias=False, - ) - self.v_proj = nn.Linear( - self.hidden_size, - self.num_kv_heads * self.head_dim, - bias=False, - ) - self.o_proj = nn.Linear( - self.num_heads * self.head_dim, - self.hidden_size, - bias=False, - ) - self.q_norm = _RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps) - self.k_norm = _RMSNorm(self.head_dim, eps=cp_cfg.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - ) -> torch.Tensor: - bsz, seq_len, _ = hidden_states.shape - hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim) - hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim) - - q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2) - k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2) - v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2) - - cos, sin = position_embeddings - # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads - cos = cos.unsqueeze(1) # [batch, 1, seq_len, head_dim] - sin = sin.unsqueeze(1) - q = (q * cos) + (_rotate_half(q) * sin) - k = (k * cos) + (_rotate_half(k) * sin) - - attn_out = F.scaled_dot_product_attention( - q, - k, - v, - scale=self.scaling, - is_causal=True, - enable_gqa=self._use_gqa, - ) - - attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) - output = self.o_proj(attn_out) - return output - - -# =================================================================== -# MLP -# =================================================================== - - -class Qwen3OmniCodePredictorMLP(nn.Module): - """SiLU-gated MLP for code predictor, matching HF's implementation.""" - - def __init__( - self, - config, - prefix: str = "", - ): - super().__init__() - hidden_size = config.code_predictor_config.hidden_size - intermediate_size = config.code_predictor_config.intermediate_size - - self.gate_proj = nn.Linear(hidden_size, intermediate_size, bias=False) - self.up_proj = nn.Linear(hidden_size, intermediate_size, bias=False) - self.down_proj = nn.Linear(intermediate_size, hidden_size, bias=False) - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - return self.down_proj(F.silu(self.gate_proj(hidden_states)) * self.up_proj(hidden_states)) - - -# =================================================================== -# Decoder Layer -# =================================================================== - - -class Qwen3OmniCodePredictorDecoderLayer(nn.Module): - """Transformer decoder layer (SDPA, no KV cache).""" - - def __init__( - self, - config, - prefix: str = "", - ) -> None: - super().__init__() - self.self_attn = Qwen3OmniCodePredictorAttention( - config, - prefix=f"{prefix}.self_attn", - ) - self.mlp = Qwen3OmniCodePredictorMLP( - config, - prefix=f"{prefix}.mlp", - ) - cp_cfg = config.code_predictor_config - self.input_layernorm = _RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps) - self.post_attention_layernorm = _RMSNorm(cp_cfg.hidden_size, eps=cp_cfg.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - ) -> torch.Tensor: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn(hidden_states, position_embeddings) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - return hidden_states - - -# =================================================================== -# Base Transformer Model (re-prefill, no KV cache) -# =================================================================== - - -class Qwen3OmniCodePredictorBaseModel(nn.Module): - """Inner transformer for code predictor. - - Signature: ``forward(inputs_embeds, position_ids) -> hidden_states`` - -- plain Tensor in, plain Tensor out (no namedtuple). - """ - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config.code_predictor_config - self.config = config - - self.codec_embedding = nn.ModuleList( - [VocabParallelEmbedding(config.vocab_size, config.hidden_size) for _ in range(config.num_code_groups - 1)] - ) - - self.layers = nn.ModuleList( - [ - Qwen3OmniCodePredictorDecoderLayer( - vllm_config.model_config.hf_config, - prefix=f"{prefix}.layers.{idx}", - ) - for idx in range(config.num_hidden_layers) - ] - ) - self.norm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = _RotaryEmbedding(config) - - def forward( - self, - inputs_embeds: torch.Tensor, - position_ids: torch.Tensor, - ) -> torch.Tensor: - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids) - for layer in self.layers: - hidden_states = layer(hidden_states, position_embeddings) - hidden_states = self.norm(hidden_states) - return hidden_states - - -# =================================================================== -# Code Predictor Wrapper (optimized re-prefill, persistent buffers) -# =================================================================== - - -class Qwen3OmniMoeTalkerCodePredictor(nn.Module): - """Optimized code predictor -- re-prefill approach, no KV cache. - - Each AR step forwards the full growing sequence (len 2 -> num_code_groups+1) - through the transformer. The extra O(T^2) FLOPs are negligible for - short sequences, and this avoids all KV-cache management overhead. - - Optimizations: - 1. Per-call embedding buffer -- avoids cross-request aliasing. - 2. Pre-allocated position_ids -- no torch.arange per step. - 3. Cached module references -- bypass ModuleList indexing. - 4. torch.compile on inner transformer. - 5. Inline sampling (top-k + top-p) -- no custom op overhead. - """ - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - - config = vllm_config.model_config.hf_config - self.config = config - self.quant_config = vllm_config.quant_config - self.prefix = prefix - - self.num_code_groups = config.code_predictor_config.num_code_groups - self._hidden_size = config.code_predictor_config.hidden_size - - self.model = Qwen3OmniCodePredictorBaseModel( + def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None: + cp_config = vllm_config.model_config.hf_config.code_predictor_config + super().__init__( vllm_config=vllm_config, + cp_config=cp_config, + wrapper_config=CodePredictorWrapperConfig( + use_cuda_graphs=False, + use_parallel_embedding=True, + use_projection=False, + return_proj_buf=True, + sampling_mode="stored", + ), + talker_hidden_size=cp_config.hidden_size, prefix=prefix, ) - - # One lm_head per residual layer (layers 1 .. G-1) - self.lm_head = nn.ModuleList( - [ - nn.Linear( - config.code_predictor_config.hidden_size, - config.code_predictor_config.vocab_size, - bias=False, - ) - for _ in range(self.num_code_groups - 1) - ] - ) - - self.set_sampling_params() - - # Lazily initialised position ids (read-only, safe to persist) - self._pos_ids: torch.Tensor | None = None - - # Cached plain-list refs (set once) - self._lm_heads: list | None = None - self._codec_embeds: list | None = None - - # Model forward (optionally compiled) - self._model_fwd: object | None = None - - def set_sampling_params(self, top_k: int = 50, top_p: float = 0.8): - """Configure sampling parameters to maintain consistency with previous implementation.""" - self._top_k = top_k - self._top_p = top_p - logger.debug(f"Sampling parameters updated: top_k={top_k}, top_p={top_p}s") - - # ------------------------------------------------------------------ - # Lazy-init helpers - # ------------------------------------------------------------------ - - def _ensure_pos_ids(self, device: torch.device) -> None: - if self._pos_ids is not None and self._pos_ids.device == device: - return - max_seq = self.num_code_groups + 1 - # [1, max_seq] for HF-style RoPE (will be expanded to [bsz, seq_len] at use) - self._pos_ids = torch.arange(max_seq, dtype=torch.long, device=device).unsqueeze(0) - - def _ensure_cached_refs(self) -> None: - if self._lm_heads is not None: - return - self._lm_heads = list(self.lm_head) - self._codec_embeds = list(self.model.codec_embedding) - - def _ensure_model_fwd(self) -> None: - if self._model_fwd is not None: - return - if current_omni_platform.supports_torch_inductor(): - # torch.compile fuses RMSNorm/RoPE in ways that lose float32 - # precision, compounding across AR steps. Use epilogue_fusion=False - # to disable the problematic fusions while still getting kernel - # fusion benefits for the linear layers and SDPA. - self._model_fwd = torch.compile( - self.model.forward, - dynamic=True, - options={ - "epilogue_fusion": False, - }, - ) - logger.info("code_predictor: torch.compile enabled (no epilogue fusion)") - else: - self._model_fwd = self.model.forward - logger.info("code_predictor: using eager mode (no torch.compile)") - - # ------------------------------------------------------------------ - # Forward -- re-prefill + inline sampling - # ------------------------------------------------------------------ - - @torch.inference_mode() - def forward( - self, - layer0_code: torch.Tensor, - layer0_embed: torch.Tensor, - last_talker_hidden: torch.Tensor, - ) -> tuple[torch.Tensor, torch.Tensor]: - """Predict residual codebooks 1..G-1 autoregressively via re-prefill. - - Args: - layer0_code: [bsz, 1] int64 - layer0_embed: [bsz, 1, hidden_size] - last_talker_hidden: [bsz, 1, hidden_size] - - Returns: - all_codes: [bsz, num_code_groups, 1] - proj_buf: [bsz, num_code_groups + 1, hidden_size] - pos 0 = last_talker_hidden (NOT a codec embed) - pos 1 = layer0_embed - pos 2.. = `codec_embedding[i](predicted_code_i)` - """ - bsz = int(layer0_code.shape[0]) - device = layer0_code.device - dtype = last_talker_hidden.dtype - num_groups = self.num_code_groups - - # Lazy init (read-only caches only) - self._ensure_pos_ids(device) - self._ensure_model_fwd() - self._ensure_cached_refs() - - # Allocate proj_buf locally each call to avoid cross-call aliasing - max_seq = num_groups + 1 - proj_buf = torch.zeros(bsz, max_seq, self._hidden_size, dtype=dtype, device=device) - pos_ids = self._pos_ids - model_fwd = self._model_fwd - lm_heads = self._lm_heads - codec_embeds = self._codec_embeds - - # Output codes - all_codes = torch.empty(bsz, num_groups, 1, dtype=torch.int64, device=device) - all_codes[:, 0] = layer0_code - - # Fill buffer positions 0 & 1 - proj_buf[:bsz, 0:1, :] = last_talker_hidden - proj_buf[:bsz, 1:2, :] = layer0_embed - - # Autoregressive loop: predict layers 1..G-1 - for step in range(1, num_groups): - seq_len = step + 1 - projected = proj_buf[:bsz, :seq_len, :] - # position_ids: [batch, seq_len] for HF-style RoPE - step_pos_ids = pos_ids[:, :seq_len].expand(bsz, -1) - - hidden_out = model_fwd(projected, step_pos_ids) - - # Inline sampling: top-k -> top-p -> softmax -> multinomial - logits = lm_heads[step - 1](hidden_out[:, -1, :]) # [bsz, vocab] - if self._top_k > 0: - topk_vals, _ = logits.topk(self._top_k, dim=-1) - logits = logits.masked_fill(logits < topk_vals[:, -1:], float("-inf")) - if self._top_p < 1.0: - sorted_logits, sorted_idx = logits.sort(dim=-1, descending=True) - cumulative_probs = F.softmax(sorted_logits, dim=-1).cumsum(dim=-1) - # Remove tokens with cumulative probability above top_p - remove_mask = cumulative_probs - F.softmax(sorted_logits, dim=-1) >= self._top_p - sorted_logits[remove_mask] = float("-inf") - logits = sorted_logits.scatter(1, sorted_idx, sorted_logits) - probs = F.softmax(logits, dim=-1) - code = torch.multinomial(probs, num_samples=1) # [bsz, 1] - - all_codes[:, step] = code - - # Embed predicted code -> next buffer position - new_embed = codec_embeds[step - 1](code) # [batch, 1, hidden_size] - proj_buf[:bsz, step + 1 : step + 2, :] = new_embed - - return all_codes, proj_buf[:bsz] - - # ------------------------------------------------------------------ - # Weight loading - # ------------------------------------------------------------------ - - def load_weights(self, weights: list[tuple[str, torch.Tensor]]) -> set[str]: - """Load weights directly (no fused projection remapping needed). - - Since we use separate nn.Linear for q/k/v/o and gate/up/down, - weight names match the HF checkpoint directly. - """ - params_dict = dict(self.named_parameters()) - loaded_params: set[str] = set() - - for name, loaded_weight in weights: - # Skip rotary embeddings - if "rotary_emb.inv_freq" in name: - continue - - param = params_dict.get(name) - if param is None: - continue - - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - - return loaded_params diff --git a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py index 1e84eaeba..8d2f0686a 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py +++ b/vllm_omni/model_executor/models/qwen3_tts/qwen3_tts_code_predictor_vllm.py @@ -1,318 +1,27 @@ +"""Qwen3-TTS Code Predictor -- thin wrapper over CodePredictorWrapper.""" + from __future__ import annotations from collections.abc import Iterable import torch -import torch.nn as nn -import torch.nn.functional as F from vllm.config import VllmConfig from vllm.config.vllm import set_current_vllm_config -from vllm.logger import init_logger -from vllm.model_executor.model_loader.weight_utils import ( - default_weight_loader, -) -from vllm_omni.platforms import current_omni_platform +from vllm_omni.model_executor.models.common.qwen3_code_predictor import ( + CodePredictorBaseModel, + CodePredictorWrapper, + CodePredictorWrapperConfig, +) from .configuration_qwen3_tts import Qwen3TTSTalkerCodePredictorConfig, Qwen3TTSTalkerConfig -logger = init_logger(__name__) +# Backward-compat alias used by tests +Qwen3TTSTalkerCodePredictorModelVLLM = CodePredictorBaseModel -# =================================================================== -# HF-numerics-compatible layers for code predictor -# =================================================================== -# -# These use plain PyTorch ops (nn.Linear, manual RMSNorm in float32, -# rotate_half RoPE) to produce outputs numerically identical to the -# HuggingFace reference. vLLM's fused kernels (RMSNorm, QKVParallel, -# get_rope) introduce small precision differences that compound across -# the 15 autoregressive steps of the code predictor, causing severe -# audio quality degradation (UTMOS ~4.26 → ~2.66). -# -# See: https://github.com/vllm-project/vllm-omni/issues/2274 - - -class _RMSNorm(nn.Module): - """RMSNorm matching HuggingFace's Qwen3TTSRMSNorm exactly. - - Computes variance in float32 to avoid bfloat16 precision loss. - """ - - def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: - super().__init__() - self.weight = nn.Parameter(torch.ones(hidden_size)) - self.variance_epsilon = eps - - def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: - input_dtype = hidden_states.dtype - hidden_states = hidden_states.to(torch.float32) - variance = hidden_states.pow(2).mean(-1, keepdim=True) - hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) - return self.weight * hidden_states.to(input_dtype) - - -def _rotate_half(x: torch.Tensor) -> torch.Tensor: - """Rotates half the hidden dims of the input.""" - x1 = x[..., : x.shape[-1] // 2] - x2 = x[..., x.shape[-1] // 2 :] - return torch.cat((-x2, x1), dim=-1) - - -class _RotaryEmbedding(nn.Module): - """RoPE matching HuggingFace's Qwen3TTSRotaryEmbedding exactly. - - Forces float32 computation for cos/sin, matching HF's torch.autocast(enabled=False). - """ - - def __init__(self, config: Qwen3TTSTalkerCodePredictorConfig) -> None: - super().__init__() - head_dim = getattr( - config, - "head_dim", - config.hidden_size // config.num_attention_heads, - ) - # Standard default RoPE - rope_theta = getattr(config, "rope_theta", 10000.0) - inv_freq = 1.0 / (rope_theta ** (torch.arange(0, head_dim, 2, dtype=torch.float32) / head_dim)) - self.register_buffer("inv_freq", inv_freq, persistent=False) - - def forward(self, x: torch.Tensor, position_ids: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: - # position_ids: [batch, seq_len] - inv_freq_expanded = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) - position_ids_expanded = position_ids[:, None, :].float() - - # Force float32 (matching HF) - device_type = x.device.type if isinstance(x.device.type, str) and x.device.type != "mps" else "cpu" - with torch.autocast(device_type=device_type, enabled=False): - freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(1, 2) - emb = torch.cat((freqs, freqs), dim=-1) - cos = emb.cos() - sin = emb.sin() - - return cos.to(dtype=x.dtype), sin.to(dtype=x.dtype) - - -class _CodePredictorAttention(nn.Module): - """Standalone multi-head attention for code predictor. - - Uses F.scaled_dot_product_attention with HF-compatible RoPE and RMSNorm. - Input: [B, seq_len, hidden_size], output: [B, seq_len, hidden_size]. - """ - - def __init__( - self, - config: Qwen3TTSTalkerCodePredictorConfig, - *, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = config.hidden_size - self.num_heads = config.num_attention_heads - self.num_kv_heads = config.num_key_value_heads - self.head_dim = getattr( - config, - "head_dim", - config.hidden_size // config.num_attention_heads, - ) - self.scaling = self.head_dim**-0.5 - self._use_gqa = self.num_kv_heads != self.num_heads - - # Separate q/k/v projections matching HF (no fused packing) - self.q_proj = nn.Linear( - self.hidden_size, - self.num_heads * self.head_dim, - bias=getattr(config, "attention_bias", False), - ) - self.k_proj = nn.Linear( - self.hidden_size, - self.num_kv_heads * self.head_dim, - bias=getattr(config, "attention_bias", False), - ) - self.v_proj = nn.Linear( - self.hidden_size, - self.num_kv_heads * self.head_dim, - bias=getattr(config, "attention_bias", False), - ) - self.o_proj = nn.Linear( - self.num_heads * self.head_dim, - self.hidden_size, - bias=False, - ) - self.q_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) - self.k_norm = _RMSNorm(self.head_dim, eps=config.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - ) -> torch.Tensor: - bsz, seq_len, _ = hidden_states.shape - hidden_shape_q = (bsz, seq_len, self.num_heads, self.head_dim) - hidden_shape_kv = (bsz, seq_len, self.num_kv_heads, self.head_dim) - - q = self.q_norm(self.q_proj(hidden_states).view(hidden_shape_q)).transpose(1, 2) - k = self.k_norm(self.k_proj(hidden_states).view(hidden_shape_kv)).transpose(1, 2) - v = self.v_proj(hidden_states).view(hidden_shape_kv).transpose(1, 2) - - cos, sin = position_embeddings - # cos/sin are [batch, seq_len, head_dim], need unsqueeze at dim=1 for heads - cos = cos.unsqueeze(1) # [batch, 1, seq_len, head_dim] - sin = sin.unsqueeze(1) - q = (q * cos) + (_rotate_half(q) * sin) - k = (k * cos) + (_rotate_half(k) * sin) - - attn_out = F.scaled_dot_product_attention( - q, - k, - v, - scale=self.scaling, - is_causal=True, - enable_gqa=self._use_gqa, - ) - - attn_out = attn_out.transpose(1, 2).reshape(bsz, seq_len, -1) - output = self.o_proj(attn_out) - return output - - -class _CodePredictorMLP(nn.Module): - """SiLU-gated MLP for code predictor, matching HF's Qwen3TTSTalkerTextMLP.""" - - def __init__( - self, - config: Qwen3TTSTalkerCodePredictorConfig, - *, - prefix: str = "", - ) -> None: - super().__init__() - self.gate_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) - self.up_proj = nn.Linear(config.hidden_size, config.intermediate_size, bias=False) - self.down_proj = nn.Linear(config.intermediate_size, config.hidden_size, bias=False) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - return self.down_proj(F.silu(self.gate_proj(x)) * self.up_proj(x)) - - -class _CodePredictorDecoderLayer(nn.Module): - """Transformer decoder layer for code predictor (SDPA, no KV cache).""" - - def __init__( - self, - config: Qwen3TTSTalkerCodePredictorConfig, - *, - prefix: str = "", - ) -> None: - super().__init__() - self.self_attn = _CodePredictorAttention(config, prefix=f"{prefix}.self_attn") - self.mlp = _CodePredictorMLP(config, prefix=f"{prefix}.mlp") - self.input_layernorm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.post_attention_layernorm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - - def forward( - self, - hidden_states: torch.Tensor, - position_embeddings: tuple[torch.Tensor, torch.Tensor], - ) -> torch.Tensor: - residual = hidden_states - hidden_states = self.input_layernorm(hidden_states) - hidden_states = self.self_attn(hidden_states, position_embeddings) - hidden_states = residual + hidden_states - - residual = hidden_states - hidden_states = self.post_attention_layernorm(hidden_states) - hidden_states = self.mlp(hidden_states) - hidden_states = residual + hidden_states - return hidden_states - - -# =================================================================== -# Code Predictor Transformer Model -# =================================================================== - - -class Qwen3TTSTalkerCodePredictorModelVLLM(nn.Module): - """Transformer model for the code predictor (re-prefill, no KV cache).""" - - def __init__( - self, - config: Qwen3TTSTalkerCodePredictorConfig, - *, - talker_hidden_size: int | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.config = config - - self.layers = nn.ModuleList( - [_CodePredictorDecoderLayer(config, prefix=f"{prefix}.layers.{i}") for i in range(config.num_hidden_layers)] - ) - self.norm = _RMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.rotary_emb = _RotaryEmbedding(config) - - # Codec embeddings: one per residual group. Stored in talker hidden dim - # (some checkpoints use talker_hidden_size != code_predictor hidden_size). - emb_dim = int(talker_hidden_size) if talker_hidden_size is not None else int(config.hidden_size) - self.codec_embedding = nn.ModuleList( - [nn.Embedding(config.vocab_size, emb_dim) for _ in range(config.num_code_groups - 1)] - ) - - def get_input_embeddings(self) -> nn.ModuleList: - return self.codec_embedding - - def forward( - self, - inputs_embeds: torch.Tensor, - position_ids: torch.Tensor, - ) -> torch.Tensor: - hidden_states = inputs_embeds - position_embeddings = self.rotary_emb(hidden_states, position_ids) - for layer in self.layers: - hidden_states = layer(hidden_states, position_embeddings) - hidden_states = self.norm(hidden_states) - return hidden_states - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - for name, loaded_weight in weights: - if "rotary_emb.inv_freq" in name: - continue - param = params_dict.get(name) - if param is None: - continue - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - - -# =================================================================== -# Code Predictor Wrapper (optimized re-prefill + torch.compile) -# =================================================================== - - -class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): - """vLLM-native code_predictor for the AR talker (residual codebooks). - - Re-prefill approach: each AR step forwards the full growing sequence - through the 5-layer transformer. No KV cache needed. This trades - ~O(T^2) extra attention FLOPs (negligible for T=16, 5 layers) for - zero KV cache management overhead and a simpler execution model. - - Uses HF-compatible layers (plain nn.Linear, float32 RMSNorm, rotate_half - RoPE) to ensure numerical fidelity with the reference implementation. - Precision matters here because small errors compound across 15 AR steps. - - Optimizations preserved: - 1. torch.compile on model forward -- fuses small kernel launches. - 2. Pre-allocated embedding buffer [B, max_seq, H] -- no torch.cat per step. - 3. Projection caching -- each token projected once and cached. - 4. Pre-allocated position_ids -- no torch.arange per step. - 5. Inline sampling -- no custom op / forward_context overhead. - 6. Cached module references -- bypass nn.Module.__call__ overhead. - 7. CUDA graphs per batch-size bucket. - """ +class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(CodePredictorWrapper): + """Qwen3-TTS code predictor (CUDA graphs, per-call sampling, projection).""" def __init__( self, @@ -322,250 +31,24 @@ class Qwen3TTSTalkerCodePredictorForConditionalGenerationVLLM(nn.Module): talker_config: Qwen3TTSTalkerConfig, prefix: str = "code_predictor", ) -> None: - super().__init__() - self._vllm_config = vllm_config - self.config = config - self.talker_config = talker_config - - self.model = Qwen3TTSTalkerCodePredictorModelVLLM( - config, + super().__init__( + vllm_config=vllm_config, + cp_config=config, + wrapper_config=CodePredictorWrapperConfig( + use_cuda_graphs=True, + use_parallel_embedding=False, + use_projection=(config.hidden_size != talker_config.hidden_size), + return_proj_buf=False, + sampling_mode="per_call", + ), talker_hidden_size=int(talker_config.hidden_size), - prefix=f"{prefix}.model", + prefix=prefix, ) - - self.lm_head = nn.ModuleList( - [nn.Linear(config.hidden_size, config.vocab_size, bias=False) for _ in range(config.num_code_groups - 1)] - ) - - if config.hidden_size != talker_config.hidden_size: - self.small_to_mtp_projection = nn.Linear(talker_config.hidden_size, config.hidden_size, bias=True) - else: - self.small_to_mtp_projection = nn.Identity() - - self._num_groups = int(config.num_code_groups) - self._talker_hidden = int(talker_config.hidden_size) - self._cp_hidden = int(config.hidden_size) - - # Pre-allocated buffers (lazily initialized on first forward). - self._proj_buf: torch.Tensor | None = None - self._model_dtype: torch.dtype | None = None - - # torch.compile + warmup state (lazily initialized in _setup_compile). - self._compiled_model_fwd = None - self._bucket_sizes: list[int] = [] - self._bucket_pos_ids: dict[int, torch.Tensor] = {} - self._lm_heads_list: list[nn.Module] | None = None - self._codec_embeds_list: list[nn.Module] | None = None - self._cuda_graphs: dict[int, tuple[torch.cuda.CUDAGraph, torch.Tensor]] = {} - - def get_input_embeddings(self) -> nn.ModuleList: - return self.model.get_input_embeddings() + # Store talker_config for backward compat (accessed by some callers) + self.talker_config = talker_config + self._vllm_config = vllm_config def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights with vllm config context (required for VocabParallelEmbedding).""" with set_current_vllm_config(self._vllm_config): - loaded: set[str] = set() - model_weights: list[tuple[str, torch.Tensor]] = [] - other_weights: list[tuple[str, torch.Tensor]] = [] - for name, w in weights: - if name.startswith("model."): - model_weights.append((name[len("model.") :], w)) - else: - other_weights.append((name, w)) - - loaded_model = self.model.load_weights(model_weights) - loaded |= {f"model.{n}" for n in loaded_model} - - params = dict(self.named_parameters(remove_duplicate=False)) - for name, w in other_weights: - if name not in params: - continue - default_weight_loader(params[name], w) - loaded.add(name) - - return loaded - - # ------------------------------------------------------------------ - # Pre-allocated buffer management - # ------------------------------------------------------------------ - - def _ensure_buffers(self, device: torch.device, dtype: torch.dtype) -> None: - max_seq = self._num_groups + 1 - if self._proj_buf is not None and self._proj_buf.device == device and self._proj_buf.dtype == dtype: - return - max_bsz = self._vllm_config.scheduler_config.max_num_seqs - self._proj_buf = torch.zeros( - max_bsz, - max_seq, - self._cp_hidden, - dtype=dtype, - device=device, - ) - - def _setup_compile(self) -> None: - """Lazily set up torch.compile with manual CUDA graph capture.""" - if self._compiled_model_fwd is not None: - return - # Cache model parameter dtype so forward() doesn't need to query it - # on every call. Also ensures warmup buffers match model precision - # even when upstream modules produce a different dtype (#2385). - self._model_dtype = next(self.model.parameters()).dtype - self._lm_heads_list = list(self.lm_head) - self._codec_embeds_list = list(self.model.codec_embedding) - if not current_omni_platform.supports_torch_inductor(): - logger.warning_once("code_predictor: torch.compile disabled") - self._compiled_model_fwd = self.model.forward - return - - # torch.compile fuses RMSNorm/RoPE in ways that lose float32 - # precision, compounding across 15 AR steps. Use torch.compile - # with options that disable the problematic fusions while still - # getting kernel fusion benefits for the linear layers and SDPA. - self._compiled_model_fwd = torch.compile( - self.model.forward, - dynamic=False, - options={ - "epilogue_fusion": False, - }, - ) - self._warmup_buckets() - self._capture_cuda_graphs() - logger.info("code_predictor: torch.compile (no epilogue fusion) + CUDA graphs") - - def _padded_bsz(self, bsz: int) -> int: - for bucket in self._bucket_sizes: - if bsz <= bucket: - return bucket - return bsz - - def _warmup_buckets(self) -> None: - """Warmup power-of-2 batch-size buckets to front-load Inductor compilation.""" - max_bsz = self._vllm_config.scheduler_config.max_num_seqs - bucket_sizes = [1 << i for i in range(max_bsz.bit_length()) if (1 << i) <= max_bsz] - if max_bsz not in bucket_sizes: - bucket_sizes.append(max_bsz) - self._bucket_sizes = sorted(bucket_sizes) - - max_seq = self._num_groups + 1 - device = next(self.model.parameters()).device - - # Ensure proj_buf matches model parameter dtype to avoid dtype - # mismatch during warmup compilation (see #2385). - self._ensure_buffers(device, self._model_dtype) - proj_buf = self._proj_buf - for bsz in self._bucket_sizes: - # position_ids: [batch, seq_len] for HF-style RoPE - pos_ids = torch.arange(max_seq, device=device, dtype=torch.long).unsqueeze(0).expand(bsz, -1) - self._bucket_pos_ids[bsz] = pos_ids - for _ in range(3): - self._compiled_model_fwd(proj_buf[:bsz, :max_seq, :], pos_ids) - logger.info("code_predictor: warmup done for buckets %s", self._bucket_sizes) - - def _capture_cuda_graphs(self) -> None: - """Capture a CUDA graph per bucket using vLLM's global graph pool.""" - from vllm.platforms import current_platform - - pool = current_platform.get_global_graph_pool() - - max_seq = self._num_groups + 1 - proj_buf = self._proj_buf - - for bsz in self._bucket_sizes: - static_input = proj_buf[:bsz, :max_seq, :] - pos_ids = self._bucket_pos_ids[bsz] - - g = torch.cuda.CUDAGraph() - with torch.cuda.graph(g, pool=pool): - static_output = self._compiled_model_fwd(static_input, pos_ids) - - self._cuda_graphs[bsz] = (g, static_output) - - logger.info("code_predictor: captured CUDA graphs for buckets %s", self._bucket_sizes) - - # ------------------------------------------------------------------ - # Optimized forward: re-prefill + torch.compile + projection cache - # ------------------------------------------------------------------ - - @torch.inference_mode() - def forward( - self, - layer0_code: torch.Tensor, - layer0_embed: torch.Tensor, - last_talker_hidden: torch.Tensor, - do_sample: bool = True, - temperature: float = 0.9, - top_k: int = 50, - top_p: float = 1.0, - ) -> torch.Tensor: - """Predict residual codebooks 1..Q-1 autoregressively via re-prefill. - - torch.compile fuses the ~60 small kernel launches per step into fewer - fused kernels, reducing kernel launch overhead by ~75%. - - Projection caching: each token is projected once via small_to_mtp_projection - and cached in _proj_buf, avoiding redundant re-projection of past tokens. - """ - bsz = int(layer0_code.shape[0]) - num_groups = self._num_groups - device = layer0_code.device - - all_codes = torch.empty(bsz, num_groups, dtype=torch.long, device=device) - all_codes[:, 0] = layer0_code.reshape(bsz) - - # _setup_compile caches _model_dtype on first call; use it for buffers - # so they always match model weight precision (#2385). - self._setup_compile() - dtype = self._model_dtype - self._ensure_buffers(device, dtype) - - proj_buf = self._proj_buf - max_seq = self._num_groups + 1 - - projection = self.small_to_mtp_projection - model_fwd = self._compiled_model_fwd - lm_heads = self._lm_heads_list - codec_embeds = self._codec_embeds_list - - use_sampling = do_sample and temperature > 0 - inv_temperature = 1.0 / max(temperature, 1e-6) if use_sampling else 0.0 - if use_sampling and top_p != 1.0: - raise NotImplementedError( - "top_p sampling is not implemented for the vLLM-native code predictor; please set top_p=1.0." - ) - - padded_bsz = self._padded_bsz(bsz) - proj_buf[:padded_bsz].zero_() - - proj_buf[:bsz, 0, :] = projection(last_talker_hidden.reshape(bsz, 1, -1).to(dtype)).reshape(bsz, -1) - proj_buf[:bsz, 1, :] = projection(layer0_embed.reshape(bsz, 1, -1).to(dtype)).reshape(bsz, -1) - full_pos_ids = self._bucket_pos_ids.get(padded_bsz) - if full_pos_ids is None: - full_pos_ids = torch.arange(max_seq, device=device, dtype=torch.long).unsqueeze(0).expand(padded_bsz, -1) - - # Use captured CUDA graph if available, otherwise call compiled fn. - cuda_graph_entry = self._cuda_graphs.get(padded_bsz) - - for step in range(1, num_groups): - if cuda_graph_entry is not None: - cuda_graph_entry[0].replay() - hidden_out = cuda_graph_entry[1] - else: - hidden_out = model_fwd(proj_buf[:padded_bsz, :max_seq, :], full_pos_ids) - logits = lm_heads[step - 1](hidden_out[:bsz, step, :]) - - if use_sampling: - scaled = logits * inv_temperature - if top_k > 0: - topk_vals, _ = scaled.topk(top_k, dim=-1) - scaled = scaled.masked_fill(scaled < topk_vals[:, -1:], float("-inf")) - probs = F.softmax(scaled, dim=-1) - next_ids = torch.multinomial(probs, num_samples=1) - else: - next_ids = logits.argmax(dim=-1, keepdim=True) - - all_codes[:, step] = next_ids.reshape(bsz) - - if step < num_groups - 1: - new_embed = codec_embeds[step - 1](next_ids) - proj_buf[:bsz, step + 1, :] = projection(new_embed.reshape(bsz, 1, -1)).reshape(bsz, -1) - - return all_codes + return super().load_weights(weights) From 50ae1de7da006324942715fd5c03d298290065de Mon Sep 17 00:00:00 2001 From: "Y. Fisher" Date: Wed, 15 Apr 2026 15:54:38 +0800 Subject: [PATCH 14/97] [Feature] HunyuanImage3 allow guidance_scale<=1 in DiT stage (#2762) Signed-off-by: KexiongYu --- .../models/hunyuan_image3/hunyuan_image3_transformer.py | 3 ++- .../models/hunyuan_image3/pipeline_hunyuan_image3.py | 6 ++---- 2 files changed, 4 insertions(+), 5 deletions(-) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py index 327260ee0..fbdacddaf 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_image3_transformer.py @@ -1684,7 +1684,8 @@ class HunYuanAttention(nn.Module): else: attn_output = self.attn(q, k, v) # For o_proj - attn_output = attn_output.view(q.shape[0], -1) + # image_attn may return a non-contiguous tensor; reshape is safe here. + attn_output = attn_output.reshape(q.shape[0], -1) output, _ = self.o_proj(attn_output) output = output.reshape(bsz, q_len, -1) return output, None, past_key_value diff --git a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py index 2f140b48f..3de0ab310 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/pipeline_hunyuan_image3.py @@ -6,7 +6,6 @@ import os from collections.abc import Iterable from typing import Any -import numpy as np import torch import torch.nn as nn from diffusers.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler @@ -544,7 +543,7 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin, Diffu generator = [torch.Generator(self.device).manual_seed(seed) for seed in seeds] # 3. apply chat template - cfg_factor = {"gen_text": 1, "gen_image": 2} + cfg_factor = {"gen_text": 1, "gen_image": 1 + int(guidance_scale > 1.0)} bot_task = kwargs.pop("bot_task", "auto") # If `drop_think` enabled, always drop parts in the context. drop_think = kwargs.get("drop_think", self.generation_config.drop_think) @@ -1009,8 +1008,7 @@ class HunyuanImage3Pipeline(HunyuanImage3PreTrainedModel, GenerationMixin, Diffu if req.sampling_params.guidance_scale_provided: guidance_scale = req.sampling_params.guidance_scale if guidance_scale <= 1.0: - logger.warning("HunyuanImage3.0 does not support guidance_scale <= 1.0, will set it to 1.0 + epsilon.") - guidance_scale = 1.0 + np.finfo(float).eps + logger.info("HunyuanImage3.0 runs without classifier-free guidance when guidance_scale <= 1.0.") image_size = (height, width) model_inputs = self.prepare_model_inputs( prompt=prompt, From c6d76d081b3e926ea44bece356889f846445440a Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Wed, 15 Apr 2026 22:25:41 +0800 Subject: [PATCH 15/97] [Bugfix] Fix broken fp8 quantisation on Z-Image-Turbo, Qwen-Image, FLUX.1-dev (#2795) Signed-off-by: Zhang Co-authored-by: pjh4993 --- .../diffusion/models/flux/flux_transformer.py | 12 ++++-- .../qwen_image/qwen_image_transformer.py | 31 ++++++++----- .../models/z_image/z_image_transformer.py | 43 ++++++++++++++++--- 3 files changed, 66 insertions(+), 20 deletions(-) diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 680b8bfbb..297c62675 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -381,7 +381,9 @@ class FluxSingleTransformerBlock(nn.Module): super().__init__() self.mlp_hidden_dim = int(dim * mlp_ratio) - self.norm = AdaLayerNormZeroSingle(dim, quant_config=quant_config, prefix=f"{prefix}.norm") + # Modulation linear kept full precision; shift/scale/gate outputs + # are multiplied into the residual stream every block (see #2728). + self.norm = AdaLayerNormZeroSingle(dim, quant_config=None, prefix=f"{prefix}.norm") self.proj_mlp = ReplicatedLinear( dim, self.mlp_hidden_dim, @@ -563,13 +565,16 @@ class FluxTransformer2DModel(nn.Module): self.context_embedder = nn.Linear(joint_attention_dim, self.inner_dim) self.x_embedder = nn.Linear(in_channels, self.inner_dim) + # Dual-stream blocks kept full precision — FP8 on their joint + # attention path causes noise on FLUX (#2728). Single-stream + # blocks (38 vs 19) still get FP8 for memory savings. self.transformer_blocks = nn.ModuleList( [ FluxTransformerBlock( dim=self.inner_dim, num_attention_heads=num_attention_heads, attention_head_dim=attention_head_dim, - quant_config=quant_config, + quant_config=None, prefix=f"transformer_blocks.{i}", ) for i in range(num_layers) @@ -589,12 +594,13 @@ class FluxTransformer2DModel(nn.Module): ] ) + # Final modulation feeds proj_out; keep full precision (see #2728). self.norm_out = AdaLayerNormContinuous( self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6, - quant_config=quant_config, + quant_config=None, prefix="norm_out", ) self.proj_out = nn.Linear(self.inner_dim, patch_size * patch_size * self.out_channels, bias=True) diff --git a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py index 9f16d8808..88a66d7f6 100644 --- a/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py +++ b/vllm_omni/diffusion/models/qwen_image/qwen_image_transformer.py @@ -169,12 +169,15 @@ class QwenTimestepProjEmbeddings(nn.Module): self.time_proj = Timesteps(num_channels=256, flip_sin_to_cos=True, downscale_freq_shift=0, scale=1000) self.timestep_embedder = TimestepEmbedding(in_channels=256, time_embed_dim=embedding_dim) + # Time embedding MLP is kept full precision (quant_config=None) — + # small layers that feed per-block modulation; precision-sensitive + # (see #2728). self.timestep_embedder.linear_1 = ReplicatedLinear( 256, embedding_dim, bias=True, return_bias=False, - quant_config=quant_config, + quant_config=None, prefix="timestep_embedder.linear_1", ) self.timestep_embedder.linear_2 = ReplicatedLinear( @@ -182,7 +185,7 @@ class QwenTimestepProjEmbeddings(nn.Module): embedding_dim, bias=True, return_bias=False, - quant_config=quant_config, + quant_config=None, prefix="timestep_embedder.linear_2", ) self.use_additional_t_cond = use_additional_t_cond @@ -701,7 +704,10 @@ class QwenImageTransformerBlock(nn.Module): self.num_attention_heads = num_attention_heads self.attention_head_dim = attention_head_dim - # Image processing modules + # Image processing modules. + # Modulation linear is kept full precision (quant_config=None) — it + # produces shift/scale/gate values that are precision-sensitive + # (see #2728). self.img_mod = nn.Sequential( nn.SiLU(), ReplicatedLinear( @@ -709,7 +715,7 @@ class QwenImageTransformerBlock(nn.Module): 6 * dim, bias=True, return_bias=False, - quant_config=quant_config, + quant_config=None, prefix="img_mod.1", ), ) @@ -725,7 +731,7 @@ class QwenImageTransformerBlock(nn.Module): self.img_norm2 = AdaLayerNorm(dim, elementwise_affine=False, eps=eps) self.img_mlp = FeedForward(dim=dim, dim_out=dim, quant_config=quant_config, prefix="img_mlp") - # Text processing modules + # Text processing modules. self.txt_mod = nn.Sequential( nn.SiLU(), ReplicatedLinear( @@ -733,7 +739,7 @@ class QwenImageTransformerBlock(nn.Module): 6 * dim, bias=True, return_bias=False, - quant_config=quant_config, + quant_config=None, prefix="txt_mod.1", ), ) @@ -963,12 +969,14 @@ class QwenImageTransformer2DModel(CachedTransformer): self.txt_norm = RMSNorm(joint_attention_dim, eps=1e-6) + # Entry projections (image/text) are kept full precision — + # small sensitive layers at the network boundary (see #2728). self.img_in = ReplicatedLinear( in_channels, self.inner_dim, bias=True, return_bias=False, - quant_config=quant_config, + quant_config=None, prefix="img_in", ) self.txt_in = ReplicatedLinear( @@ -976,7 +984,7 @@ class QwenImageTransformer2DModel(CachedTransformer): self.inner_dim, bias=True, return_bias=False, - quant_config=quant_config, + quant_config=None, prefix="txt_in", ) @@ -993,13 +1001,16 @@ class QwenImageTransformer2DModel(CachedTransformer): ] ) + # Final modulation and output projection are kept full precision — + # they produce the output latent and are precision-sensitive + # (see #2728). self.norm_out = AdaLayerNormContinuous(self.inner_dim, self.inner_dim, elementwise_affine=False, eps=1e-6) self.norm_out.linear = ReplicatedLinear( self.inner_dim, 2 * self.inner_dim, bias=True, return_bias=False, - quant_config=quant_config, + quant_config=None, prefix="norm_out.linear", ) self.proj_out = ReplicatedLinear( @@ -1007,7 +1018,7 @@ class QwenImageTransformer2DModel(CachedTransformer): patch_size * patch_size * self.out_channels, bias=True, return_bias=False, - quant_config=quant_config, + quant_config=None, prefix="proj_out", ) diff --git a/vllm_omni/diffusion/models/z_image/z_image_transformer.py b/vllm_omni/diffusion/models/z_image/z_image_transformer.py index 3ffad221b..c36ea7466 100644 --- a/vllm_omni/diffusion/models/z_image/z_image_transformer.py +++ b/vllm_omni/diffusion/models/z_image/z_image_transformer.py @@ -214,12 +214,14 @@ class TimestepEmbedder(nn.Module): super().__init__() if mid_size is None: mid_size = out_size + # Time embedding MLP is kept full precision (quant_config=None) — + # small layers that feed adaLN; precision-sensitive (see #2728). self.mlp = nn.Sequential( ReplicatedLinear( frequency_embedding_size, mid_size, bias=True, - quant_config=quant_config, + quant_config=None, return_bias=False, ), nn.SiLU(), @@ -227,7 +229,7 @@ class TimestepEmbedder(nn.Module): mid_size, out_size, bias=True, - quant_config=quant_config, + quant_config=None, return_bias=False, ), ) @@ -426,9 +428,16 @@ class ZImageTransformerBlock(nn.Module): self.modulation = modulation if modulation: + # Modulation linear is kept at full precision (quant_config=None) + # — it produces scale/gate values that are precision-sensitive + # (see #2728, mirrors OmniGen2 fix). self.adaLN_modulation = nn.Sequential( ReplicatedLinear( - min(dim, ADALN_EMBED_DIM), 4 * dim, bias=True, return_bias=False, quant_config=quant_config + min(dim, ADALN_EMBED_DIM), + 4 * dim, + bias=True, + quant_config=None, + return_bias=False, ), ) @@ -485,14 +494,24 @@ class FinalLayer(nn.Module): def __init__(self, hidden_size, out_channels, quant_config: "QuantizationConfig | None" = None): super().__init__() self.norm_final = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6) + # Final output projection and its modulation are precision-sensitive + # (produce the output latent); keep at full precision (see #2728). self.linear = ReplicatedLinear( - hidden_size, out_channels, bias=True, quant_config=quant_config, return_bias=False + hidden_size, + out_channels, + bias=True, + quant_config=None, + return_bias=False, ) self.adaLN_modulation = nn.Sequential( nn.SiLU(), ReplicatedLinear( - min(hidden_size, ADALN_EMBED_DIM), hidden_size, bias=True, quant_config=quant_config, return_bias=False + min(hidden_size, ADALN_EMBED_DIM), + hidden_size, + bias=True, + quant_config=None, + return_bias=False, ), ) @@ -673,11 +692,13 @@ class ZImageTransformer2DModel(CachedTransformer): all_x_embedder = {} all_final_layer = {} for patch_idx, (patch_size, f_patch_size) in enumerate(zip(all_patch_size, all_f_patch_size)): + # x_embedder (patch embed) is a small precision-sensitive entry + # layer; keep full precision (see #2728). x_embedder = ReplicatedLinear( f_patch_size * patch_size * patch_size * in_channels, dim, bias=True, - quant_config=quant_config, + quant_config=None, return_bias=False, ) all_x_embedder[f"{patch_size}-{f_patch_size}"] = x_embedder @@ -720,9 +741,17 @@ class ZImageTransformer2DModel(CachedTransformer): ] ) self.t_embedder = TimestepEmbedder(min(dim, ADALN_EMBED_DIM), mid_size=1024, quant_config=quant_config) + # Caption embedder maps text features -> hidden; keep full precision + # (see #2728). self.cap_embedder = nn.Sequential( RMSNorm(cap_feat_dim, eps=norm_eps), - ReplicatedLinear(cap_feat_dim, dim, bias=True, return_bias=False, quant_config=quant_config), + ReplicatedLinear( + cap_feat_dim, + dim, + bias=True, + quant_config=None, + return_bias=False, + ), ) self.x_pad_token = nn.Parameter(torch.empty((1, dim))) From f1e3f037265852b952cef654c489182bf7c26686 Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Wed, 15 Apr 2026 11:01:45 -0600 Subject: [PATCH 16/97] [feature] Hidden State Prefix Caching (#2164) Signed-off-by: Alex Brooks --- docs/.nav.yml | 1 + docs/design/feature/prefix_caching.md | 164 +++++++++ tests/conftest.py | 6 + tests/core/test_prefix_cache.py | 347 ++++++++++++++++++++ tests/e2e/online_serving/test_qwen3_omni.py | 75 ++++- vllm_omni/core/prefix_cache.py | 264 +++++++++++++++ vllm_omni/utils/mm_outputs.py | 93 ++++++ vllm_omni/worker/gpu_ar_model_runner.py | 204 +++++++++--- vllm_omni/worker/gpu_model_runner.py | 59 +++- 9 files changed, 1145 insertions(+), 68 deletions(-) create mode 100644 docs/design/feature/prefix_caching.md create mode 100644 tests/core/test_prefix_cache.py create mode 100644 vllm_omni/core/prefix_cache.py create mode 100644 vllm_omni/utils/mm_outputs.py diff --git a/docs/.nav.yml b/docs/.nav.yml index 441ef9f52..79d7c38e2 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -98,6 +98,7 @@ nav: - 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 diff --git a/docs/design/feature/prefix_caching.md b/docs/design/feature/prefix_caching.md new file mode 100644 index 000000000..ebad8b691 --- /dev/null +++ b/docs/design/feature/prefix_caching.md @@ -0,0 +1,164 @@ +# Automatic Prefix Caching in Omni Models + + +--- + +## Table of Contents + +- [Overview](#overview) +- [High-Level Approach](#high-level-approach) +- [Example](#example) +- [What About Multimodal Inputs?](#what-about-multimodal-inputs) + +--- + +### Overview + +Prefix caching in the context of kv-cache management is a useful optimization for avoiding redundant computations. The main idea is that we store portions of the kv-cache from processed requests, so that we can reuse them if incoming requests have the same prefix as previous requests. + +vLLM manages the kv-cache as blocks, which represent a span of tokens of a fixed length. Blocks are hashable by the content that they contain, which typically means the tokens within the span, but also could be influenced by other factors, e.g., LoRA and multimodal data. + +vLLM implements automatic prefix caching for managing its kv-cache, which is best understood by reading the design document [here](https://docs.vllm.ai/en/latest/design/prefix_caching/). vLLM-Omni builds on top of the prefix caching mechanism in a noninvasive way to allow caching between stages in Omni pipelines. This typically means for a given stage we aim to support caching for the following: + +- The last hidden states produced by the stage +- Model / stage specific multimodal data + +!!! note "Note 1" + This document describes vLLM-Omni's mechanism for caching tensor outputs that are meant to be passed between stages, when requests have common prefixes, similar to the way in which vLLM has prefix caching for the kv-cache. This works in conjunction with vLLM's multimodal encoder caching, but is distinct. See the final section for a concrete example for how they tie together in practice. + +### High-Level Approach +!!! note "Note 2" + Prior to reading this section, it's recommended to take a look at the design documents in vLLM for [Automatic Prefix Caching](https://docs.vllm.ai/en/latest/features/automatic_prefix_caching/), which will make some of the concepts more clear. + +The main focus of vLLM-Omni's approach to prefix caching stage outputs is to build on vLLM's prefix caching in the least invasive way possible while minimizing impact for cache misses, and consuming a minimal amount of GPU memory. To understand the implementation, there are a few important things to note: + +- Between stages, device tensors are generally moved to CPU; this is important since we're just caching the outputs of stages, so it is okay to keep the entire cache on the CPU. + +- For a tensor to be considered cacheable, the first dimension (currently) needs to be the same as the token count, as it allows us to reuse block/slot mappings for our externally maintained tensor caches. This allows us to dynamically discover the tensors to be marked as cacheable outputs in each Omni model without having to explicitly specify cacheable output field names in every model. + +With this in mind, consider the set of blocks in a 2D layout, where the row represents the index of blocks being considered, and the columns represent the slots corresponding to tokens within each block. Since we know the `num_blocks` and `block_size` from our kv cache config, if we want to cache a tensor with feature size `D`, we can preallocate a CPU tensor of size `(num_blocks, block_size, D)`, and use the same block index and slot mapping to retrieve the corresponding feature vector. + + +### Example +!!! note "Note 3" + Prefix caching in vLLM-Omni currently is only supported on AutoRegressive stages with one kv-cache group. It can be enabled/disabled per-stage via the `enable_prefix_caching` parameter in the model's stage config. + +The way in which vLLM-Omni ties into vLLM's prefix caching is best understood by example. Say that we have the following: + +- `num_blocks=8` +- `block_size=4` +- `hidden_size=2` +- A stage specific multimodal output tensor named `mm_feature` with feature dimension `16` + +The prefix cache flow is then outlined below. + +1. When the model is initialized, we can determine the `hidden_size` from the `ModelConfig`, and allocate a cache of size `(num_blocks, block_size, hidden_size)`. + +2. Say we process the request `The quick brown fox was tired and slept beneath the shady tree`, which is 12 tokens and evenly divides into 3 blocks as shown below. + +``` + [ The quick brown fox ] [ was tired and slept ] [beneath the shady tree ] +Block 1: |<--- block tokens ---->| +Block 2: |<------- prefix ------>| |<--- block tokens --->| +Block 3: |<------------------ prefix -------------------->| |<--- block tokens ---->| +``` + +When the request processes, we inspect the multimodal outputs and identify the `mm_feature` tensor, which will be of shape `(seq_len, feature_dim)`, i.e., `(12, 16)` in this example. We note that the first axis is dependent on the `seq_len` and add a new cache_tensor of shape `(num_blocks, block_size, feature_dim)` to our multimodal cache for tensors. + + +3. If we lay out the cache as a 2D tensor of shape (`num_blocks`, `block_size`), we'll have something like the following: + +``` +0: [ The quick brown fox ] +1: [ was tired and slept ] +2: [beneath the shady tree ] +3: [EMPTY] +... +7: [EMPTY] +``` + +Or, if we flatten it down to 1D, +``` +0: The +1: quick +2: brown +3: fox +... +11: tree +12: [EMPTY] +... +``` + +which we can think of as row indices into the hidden states tensor if we view it as the 2D shape `(num_blocks x block_size, feature_dim)`. That is, the analogous flattened (from 3D -> 2D) mapping of the cache for hidden states becomes the following. +``` +0: +1: +2: +3: +... +11: +12: [EMPTY] +... +``` + +Similarly, for the multimodal outputs cache, the flattened coordinates are the same, but the `mm_feature` maps to vectors of length `16` instead of the hidden size of `2`. Note that in practice, we may have multiple multimodal output tensors per forward pass, which may have different names and different feature dimensions. + + +4. Now, say that we receive a new request `The quick brown fox jumped over the dog`. + +``` + [ The quick brown fox ] [ jumped over the dog ] +Block 1: |<--- block tokens ---->| +Block 2: |<------- prefix ------>| |<--- block tokens --->| +``` + +Here, we will have a cache hit for `Block 1` which will be detected by vLLM based on the hash of the first block when it's handling the prefix caching on the kv-cache. As a result, when we get the output from the scheduler, we will see that `num_computed_tokens=4` (corresponding to the cached first block), and we only need to process the remaining 4 new tokens in the new prefill. + +Since we have the block indices / slot mappings from the kv cache manager, we can simply mirror the mappings and leverage the same indices for the cached hidden states and multimodal outputs. This allows us to look up the correct tensors from our externally maintained 3D caches. + +``` +0: [ The quick brown fox ] < already in the cache +1: [ was tired and slept ] +2: [beneath the shady tree ] +3: [ jumped over the dog ] < added on the second request +4: [EMPTY] +... +7: [EMPTY] +... +``` + +Finally, to pass the full hidden states and multimodal outputs to the next stage, we simply concatenate the cached contents with the corresponding new tensors computed from the current forward call. + + +### What About Multimodal Inputs? +It's also useful to consider the case about how Omni prefix caching is handled when we have multimodal inputs that don't cleanly end on block boundaries, as well as how this works with multimodal encoder caching in vLLM. For example: + +``` + [ Im0 Im1 Im2 Im3 ] [ Im4 Im5 foo ] +Block 1: |<--- block tokens ---->| +Block 2: |<------- prefix ------>| |<--- block tokens --->| +``` + +In this case, only `Block 1` will have outputs stored in the prefix tensor cache, because vLLM does not store partial blocks. This may appear to be a problem at first glance, because the multimodal input is fragmented across a new block that wasn't cached. + +In reality, this isn't a big problem for correctness, because vLLM also maintains an encoder cache for multimodal inputs. In other words, after the first pass, we'll have the following: + +- The Block 1 hash, which is used for prefix caching +- The hash describing the image data starting at position 0 and with length 6 +- In vLLM's encoder cache, a mapping from the image hash above to the encoder output + + +To understand what happens, say we get the following input as a second request: +``` + [ Im0 Im1 Im2 Im3 ] [ Im4 Im5 bar baz ] +Block 1: |<--- block tokens ---->| +Block 2: |<------- prefix ------>| |<--- block tokens --->| +``` + +First, the scheduler will check for a prefix cache hit, which we will see on `Block 1`. As a result, we will have 4 tokens marked as precomputed, and only see the remaining 4 tokens in the following prefill. + +Because we have multimodal data in a scheduled span that isn't fully precomputed, we still need to call the visual encoder. However, since we have the image hash and encoder cache, we will retrieve the encoder outputs for `Im4` and `Im5` as we create the multimodal embeddings. + +When we pass our multimodal tensors to the language model component in the same stage, we'll then expect the same outputs, because the prefix caching behaviors in vLLM-Omni / vLLM match, so the LLM will use vLLM's KV cache manager's prefix caching to correctly handle the attention information for `Block 1` while calculating the outputs for `Block 2`, giving us the correct results for processing `Block 2` with the context of `Block 1`. + +Finally, we look up the output hidden states/multimodal tensors corresponding to the prefix cache hit `Block 1` and concatenate it with the forward pass result to get the final result, which is expected to be identical to the full hidden states when prefix caching is disabled. diff --git a/tests/conftest.py b/tests/conftest.py index 098fd8d97..ad1008b72 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1850,6 +1850,7 @@ class OmniResponse: e2e_latency: float | None = None success: bool = False error_message: str | None = None + cached_tokens: int | None = None @dataclass @@ -2345,6 +2346,11 @@ class OpenAIClientHandler: if hasattr(choice.message, "content") and choice.message.content is not None: text_content = choice.message.content + # Extract cached_tokens for prefix caching tests + usage = getattr(chat_completion, "usage", None) + if usage and (details := getattr(usage, "prompt_tokens_details", None)): + result.cached_tokens = details.cached_tokens + # Calculate end-to-end latency result.e2e_latency = time.perf_counter() - start_time diff --git a/tests/core/test_prefix_cache.py b/tests/core/test_prefix_cache.py new file mode 100644 index 000000000..c3d8c1ff9 --- /dev/null +++ b/tests/core/test_prefix_cache.py @@ -0,0 +1,347 @@ +from unittest.mock import Mock, patch + +import pytest +import torch + +from vllm_omni.core.prefix_cache import OmniTensorPrefixCache + +DEFAULT_SEQ_LEN = 15 +NUM_BLOCKS = 10 +BLOCK_SIZE = 4 +HIDDEN_SIZE = 2 +DTYPE = torch.float32 +OTHER_DTYPE = torch.float16 +DEFAULT_SHAPE = torch.Size([NUM_BLOCKS, BLOCK_SIZE, HIDDEN_SIZE]) + + +class MockInputBatch: + def __init__(self, num_computed_tokens_cpu): + self.req_ids = ["req1", "req2"] + self.req_id_to_index = {req_id: i for i, req_id in enumerate(self.req_ids)} + self.num_computed_tokens_cpu = num_computed_tokens_cpu + # Block table is only mocked for validation of length; + # we don't actually need to add valid values here since + # we patch the table when testing. + self.block_table = Mock() + self.block_table.block_tables = [None] + + +def get_omni_pcache_with_mm_tensors(feat_dims, seq_len) -> OmniTensorPrefixCache: + """Build an OmniTensorPrefixCache and init mm tensors.""" + cache = get_omni_pcache() + mm_outputs = get_multimodal_outputs(feat_dims, seq_len) + cache.maybe_init_missing_mm_cache_keys(mm_outputs, seq_len) + return cache + + +def get_omni_pcache() -> OmniTensorPrefixCache: + """Build an OmniTensorPrefixCache, but don't init mm tensors.""" + cache = OmniTensorPrefixCache( + num_blocks=NUM_BLOCKS, + block_size=BLOCK_SIZE, + hidden_size=HIDDEN_SIZE, + hs_dtype=DTYPE, + ) + return cache + + +def get_multimodal_outputs(feat_dims: dict[str, int], seq_len: int) -> dict[str, torch.Tensor]: + fake_mm_inputs = {} + for mm_key, feat_dim in feat_dims.items(): + fake_mm_inputs[mm_key] = torch.rand((seq_len, feat_dim), dtype=DTYPE) + return fake_mm_inputs + + +### Tests for initialization +def test_initialization_simple(): + """Check default initialization only creates the hidden states.""" + cache = get_omni_pcache() + assert isinstance(cache.hidden_states_cache, torch.Tensor) + assert cache.hidden_states_cache.shape == DEFAULT_SHAPE + assert len(cache.mm_outputs_cache) == 0 + assert len(cache.mm_cache_keys) == 0 + + +def test_initialization_with_multimodal(): + """Check initialization + registration of multimodal outputs.""" + cache = get_omni_pcache() + feat_dims = {"foo": 100, "bar": 50, "baz": 10} + mm_outputs = get_multimodal_outputs( + feat_dims, + seq_len=DEFAULT_SEQ_LEN, + ) + # Cast one of the keys to a different dtype; the dtype of the tensor + # that is used to initialize the cache dictates the cache dtype. + mm_outputs["foo"] = mm_outputs["foo"].to(OTHER_DTYPE) + + cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN) + assert len(cache.mm_cache_keys) == 3 + assert set(cache.mm_cache_keys) == set(feat_dims.keys()) + for mm_key in cache.mm_cache_keys: + cache_tensor = cache.mm_outputs_cache[mm_key] + assert isinstance(cache_tensor, torch.Tensor) + assert cache_tensor.shape[-1] == feat_dims[mm_key] + assert mm_outputs[mm_key].dtype == cache_tensor.dtype + + +def test_init_missing_mm_cache_keys_is_idempotent(): + """Ensure that the cache doesn't reinitialize old keys.""" + cache = get_omni_pcache() + mm_key = "foo" + feat_dims = {mm_key: 100} + mm_outputs = get_multimodal_outputs( + feat_dims, + seq_len=DEFAULT_SEQ_LEN, + ) + cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN) + assert len(cache.mm_cache_keys) == 1 + assert mm_key in cache.mm_cache_keys + + # Cache is initialized to 0 - fill it with 1s + cache.mm_outputs_cache[mm_key].fill_(1) + + # Ensure that running another initialization + # doesn't zero out our cache values + cache.maybe_init_missing_mm_cache_keys(mm_outputs, DEFAULT_SEQ_LEN) + assert len(cache.mm_cache_keys) == 1 + assert mm_key in cache.mm_cache_keys + assert torch.all(cache.mm_outputs_cache[mm_key] == 1) + + +### Tests for Update +def test_update_no_multimodal(): + """Test that slot mappings act as row indices hidden states.""" + cache = get_omni_pcache() + + num_tokens_unpadded = 8 + slot_offset = 8 + slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded) + new_hidden_states = torch.rand((num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE) + + cache.update_omni_tensor_prefix_cache( + hidden_states=new_hidden_states, + multimodal_outputs=None, + num_tokens_unpadded=num_tokens_unpadded, + slot_mapping=slot_mapping, + ) + + # Ensure that if we reshape our 3D cache back to 2D, we can use the + # indices in our slot mappings to access the hidden states as expected + hs_rows = cache.hidden_states_cache.view(NUM_BLOCKS * BLOCK_SIZE, HIDDEN_SIZE) + for slot_idx, new_states in zip(slot_mapping, new_hidden_states): + slot_states = hs_rows[slot_idx] + assert torch.all(slot_states == new_states) + + +@pytest.mark.parametrize( + "feat_dims", + [ + {"foo": 100, "bar": 100}, + {"foo": 100, "bar": 50, "baz": 10}, + ], +) +def test_update_with_multimodal_outputs(feat_dims): + """Test that slot mappings are correct for multimodal tensors.""" + cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN) + + num_tokens_unpadded = 8 + slot_offset = 8 + slot_mapping = torch.arange(slot_offset, slot_offset + num_tokens_unpadded) + feature_dims = {key: val.shape[-1] for key, val in cache.mm_outputs_cache.items()} + mm_outputs = {key: torch.rand((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=mm_outputs, + num_tokens_unpadded=num_tokens_unpadded, + slot_mapping=slot_mapping, + ) + + for mm_key in feat_dims.keys(): + assert mm_key in cache.mm_outputs_cache + key_feat_dim = feature_dims[mm_key] + mm_state_rows = cache.mm_outputs_cache[mm_key].view(NUM_BLOCKS * BLOCK_SIZE, key_feat_dim) + + # Similar to hidden states, but for each key in the dict; + # Different tensors may have different feature dims + new_mm_outputs = mm_outputs[mm_key] + for slot_idx, new_output in zip(slot_mapping, new_mm_outputs): + slot_states = mm_state_rows[slot_idx] + assert torch.all(slot_states == new_output) + + +### Tests for Merging +def fake_get_cached_block_ids(self, req_idx, *args, **kwargs): + """Fake block table lookup. + + Assumption: + req_idx 0 is a cache hit with slots 8, 9, ..., 15 + req_idx 1 is a cache miss + """ + assert req_idx < 2 + if req_idx == 0: + # With the slot offset we provided (8), the corresponding + # blocks IDs are 2 & 3 because the block size is 4. + return torch.tensor([2, 3], dtype=torch.long) + return torch.tensor([], dtype=torch.long) + + +@pytest.mark.parametrize("num_tokens_padded", [None, 16]) +def test_get_merged_hidden_states(num_tokens_padded): + """Ensure that hidden states are merged correctly.""" + cache = get_omni_pcache() + + orig_num_tokens_unpadded = 8 + slot_offset = 8 # We'll put our states in slots 8, 9, 10, ..., 15 + orig_slot_mapping = torch.arange(slot_offset, slot_offset + orig_num_tokens_unpadded) + orig_hidden_states = torch.rand((orig_num_tokens_unpadded, HIDDEN_SIZE), dtype=DTYPE) + + cache.update_omni_tensor_prefix_cache( + hidden_states=orig_hidden_states, + multimodal_outputs=None, + num_tokens_unpadded=orig_num_tokens_unpadded, + slot_mapping=orig_slot_mapping, + num_tokens_padded=num_tokens_padded, + ) + + # Say that we have two requests, but only one of them is a cache hit + 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_hidden_states = torch.rand( + (num_new_toks_req1 + num_new_toks_req2, HIDDEN_SIZE), + dtype=DTYPE, + ) + req1_new_states = new_hidden_states[:num_new_toks_req1] + req2_new_states = new_hidden_states[-num_new_toks_req2:] + + input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])) + + with patch( + "vllm_omni.core.prefix_cache.OmniTensorPrefixCache._get_cached_block_ids", + new=fake_get_cached_block_ids, + ): + merged_states = cache.get_merged_hidden_states( + query_start_loc=[0, num_new_toks_req1], + input_batch=input_batch, + hidden_states=new_hidden_states, + num_scheduled_tokens=num_scheduled_tokens, + ) + + assert "req1" in merged_states and "req2" in merged_states + req1_merged_states = merged_states["req1"] + req2_merged_states = merged_states["req2"] + + # First, check the cache hit case + assert req1_merged_states.shape == torch.Size([orig_num_tokens_unpadded + num_new_toks_req1, HIDDEN_SIZE]) + # Ensure that the req1 merged states are the cached states + the new req1 states + assert torch.all(req1_merged_states[:orig_num_tokens_unpadded] == orig_hidden_states) + assert torch.all(req1_merged_states[-num_new_toks_req1:] == req1_new_states) + + # Next, ensure that the cache miss case only has the new states + assert req2_merged_states.shape == torch.Size([num_new_toks_req2, HIDDEN_SIZE]) + assert torch.all(req2_merged_states == req2_new_states) + + +@pytest.mark.parametrize("num_tokens_padded", [None, 16]) +@pytest.mark.parametrize( + "feat_dims", + [ + {"foo": 100, "bar": 100}, + {"foo": 100, "bar": 50, "baz": 10}, + ], +) +def test_get_merged_multimodal_outputs(feat_dims, num_tokens_padded): + cache = get_omni_pcache_with_mm_tensors(feat_dims, seq_len=DEFAULT_SEQ_LEN) + + orig_num_tokens_unpadded = 8 + slot_offset = 8 # We'll put our states in slots 8, 9, 10, ..., 15 + 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=num_tokens_padded, + ) + + # Similar to hs test- say that we have two requests, but only one of them is a cache hit + 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, + ) + # We also want to make sure passthrough data (outside of our keys) isn't dropped + new_mm_outputs["passthrough_data"] = "Something else" + # Lists are a special case because we can't split them yet if we want to match + # the nonprefix cache behavior, because this runs before post process. + new_mm_outputs["passthrough_list"] = ["should", "not", "split"] + + input_batch = MockInputBatch(num_computed_tokens_cpu=torch.Tensor([orig_num_tokens_unpadded, 0])) + + with 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, + ) + + # Ensure the passthrough data wasn't dropped + assert "passthrough_data" in merged_mm_outputs + assert "passthrough_list" in merged_mm_outputs + + for mm_key, mm_output in merged_mm_outputs.items(): + # Ensure passthrough data is just forwarded normally and not duplicated + assert isinstance(mm_output, dict) + assert "req1" in mm_output and "req2" in mm_output + if mm_key == "passthrough_data": + assert mm_key not in cache.mm_cache_keys + assert new_mm_outputs[mm_key] == mm_output["req1"] + assert new_mm_outputs[mm_key] == mm_output["req2"] + elif mm_key == "passthrough_list": + assert mm_key not in cache.mm_cache_keys + assert new_mm_outputs[mm_key] == mm_output["req1"] + assert new_mm_outputs[mm_key] == mm_output["req2"] + else: + assert mm_key in cache.mm_cache_keys + curr_feat_dim = feature_dims[mm_key] + # Ensure that req1 (cache hit) merged the mm data + req1_merged_mm_outputs = mm_output["req1"] + req1_new_mm_outputs = new_mm_outputs[mm_key][:num_new_toks_req1] + + assert req1_merged_mm_outputs.shape == torch.Size( + [orig_num_tokens_unpadded + num_new_toks_req1, curr_feat_dim] + ) + # Ensure that the req1 merged mm data are the cached data + the new data + assert torch.all(req1_merged_mm_outputs[:orig_num_tokens_unpadded] == orig_mm_outputs[mm_key]) + assert torch.all(req1_merged_mm_outputs[-num_new_toks_req1:] == req1_new_mm_outputs) + + # Ensure that req2 (cache miss) only has the new mm data + req2_merged_mm_outputs = mm_output["req2"] + req2_new_mm_outputs = new_mm_outputs[mm_key][-num_new_toks_req2:] + + 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) diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index f4aabb8b9..c05f8f506 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -23,11 +23,13 @@ os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] +QWEN3_OMNI_CONFIG_PATH = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml") +QWEN3_OMNI_XPU_CONFIG_PATH = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml") -def get_chunk_config(): +def get_chunk_config(config_path: str): path = modify_stage_config( - str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml"), + config_path, updates={ "async_chunk": True, "stage_args": { @@ -44,15 +46,41 @@ def get_chunk_config(): return path +def get_prefix_caching_config(config_path: str): + """Create a stage config with prefix caching enabled on the thinker (stage 0).""" + path = modify_stage_config( + config_path, + updates={ + "stage_args": { + 0: {"engine_args.enable_prefix_caching": True}, + }, + }, + ) + return path + + if current_omni_platform.is_xpu(): - stage_configs = [str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml")] + stage_configs = [QWEN3_OMNI_XPU_CONFIG_PATH] + prefix_caching_stage_configs = [get_prefix_caching_config(QWEN3_OMNI_XPU_CONFIG_PATH)] else: # MI325 GPU should share the same config as H100 - stage_configs = [get_chunk_config()] + stage_configs = [get_chunk_config(QWEN3_OMNI_CONFIG_PATH)] + prefix_caching_stage_configs = [get_prefix_caching_config(QWEN3_OMNI_CONFIG_PATH)] # Create parameter combinations for model and stage config test_params = [ OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs ] +# For prefix caching, we need to enable prompt token details so that we +# can determine if any tokens were cached. +prefix_test_params = [ + OmniServerParams( + model=model, + stage_config_path=stage_config, + server_args=["--enable-prompt-tokens-details"], # Enable prompt tokens details to get cached_tokens + ) + for model in models + for stage_config in prefix_caching_stage_configs +] def get_system_prompt(): @@ -75,6 +103,7 @@ def get_prompt(prompt_type="text_only"): prompts = { "text_only": "What is the capital of China? Answer in 20 words.", "mix": "What is recited in the audio? What is in this image? Describe the video briefly.", + "text_image": "What color are the squares in this image?", } return prompts.get(prompt_type, prompts["text_only"]) @@ -147,3 +176,41 @@ def test_text_to_text_001(omni_server, openai_client) -> None: } openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + + +@pytest.mark.advanced_model +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.parametrize("omni_server", prefix_test_params, indirect=True) +def test_thinker_prefix_caching(omni_server, openai_client) -> None: + """ + Test thinker prefix caching by sending identical requests with an image (i.e., + a large shared prefix) and verifying that the second request uses cached tokens + & produces the same output. + """ + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + image_data_url=image_data_url, + content_text=get_prompt("text_image"), + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": False, + "modalities": ["text"], + } + + response_1 = openai_client.send_omni_request(request_config, request_num=1)[0] + response_2 = openai_client.send_omni_request(request_config, request_num=1)[0] + + assert response_1.success + assert response_2.success + assert response_2.cached_tokens is not None + # We should cache the vast majority of the prompt (image + up to last full block), + # and set seed in the CI config, so the second request should give an identical + # response for the generated input image, even if we use dummy weights + assert response_2.cached_tokens > 0 + assert response_1.text_content == response_2.text_content diff --git a/vllm_omni/core/prefix_cache.py b/vllm_omni/core/prefix_cache.py new file mode 100644 index 000000000..69e7346c4 --- /dev/null +++ b/vllm_omni/core/prefix_cache.py @@ -0,0 +1,264 @@ +""" +Utilities for Prefix Caching in Omni models. +""" + +import torch +from vllm.logger import init_logger +from vllm.v1.worker.gpu_input_batch import InputBatch + +from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element + +logger = init_logger(__name__) + + +class OmniTensorPrefixCache: + """Prefix cache for hidden states (model outputs) and model specific + multimodal outputs. + + This class implements prefix caching in a non-invasive way on top of + vLLM by leveraging the same slot mappings that the vLLM scheduler uses + for the KV Cache. + + Conceptually, this means we are mapping vLLM's cache mapping: + (num_blocks, block_size) + + to 3D tensors of shape: + (num_blocks, block_size, feature_size) + + Note that feature_size may vary across multimodal_outputs. + """ + + def __init__( + self, + num_blocks: int, + block_size: int, + hidden_size: int, + hs_dtype: torch.dtype, + ): + self.num_blocks = num_blocks + self.block_size = block_size + self.default_hidden_size = hidden_size + + # Initialize the hidden states cache immediately + self.hidden_states_cache = self._get_cache_tensor(dtype=hs_dtype) + + # Defer initialization of the mm_outputs_cache until we + # actually see mm output tensors dependent on num tokens. + self.mm_outputs_cache = {} + self.mm_cache_keys = set() + self._new_req_cache_hit_ids: set[str] = set() + + def maybe_init_missing_mm_cache_keys(self, multimodal_outputs: dict, seq_len: int): + """Given multimodal outputs from executing the model, dynamically + determine which multimodal outputs are tensors depending on sequence + length and should be cached, and initialize the cache tensors + accordingly. + + NOTE: This is done to avoid the need for explicit specification of + cache keys for every model/stage and aligns with the current way + that we slice the multimodal outputs based on the first dimension. + + This will usually be called by the first forward pass, i.e., + determined by the warmup. + """ + for key, val in multimodal_outputs.items(): + if isinstance(val, torch.Tensor) and val.shape[0] == seq_len and key not in self.mm_cache_keys: + feat_dim = val.shape[-1] + self.mm_outputs_cache[key] = self._get_cache_tensor( + dtype=val.dtype, + hidden_size=feat_dim, + ) + self.mm_cache_keys.add(key) + new_tensor_shape = self.mm_outputs_cache[key].shape + logger.info("Initializing multimodal output cache of size %s for key: %s", list(new_tensor_shape), key) + + def _get_cache_tensor(self, dtype: torch.dtype, hidden_size: int | None = None) -> torch.Tensor: + """Allocate a CPU cache tensor for a specific key.""" + actual_hidden_size = hidden_size if hidden_size is not None else self.default_hidden_size + return torch.zeros( + (self.num_blocks, self.block_size, actual_hidden_size), + dtype=dtype, + device="cpu", + ) + + def add_prefix_cached_new_req_id(self, req_id: str): + """Adds a new request ID to the set of prefix cache hits on the batch.""" + self._new_req_cache_hit_ids.add(req_id) + + def reset_prefix_cached_new_req_ids(self): + """Clears the cache hit IDs to prepare for a new engine step.""" + self._new_req_cache_hit_ids.clear() + + @staticmethod + def _coerce_to_cpu_tensor(maybe_gpu_tensor: torch.Tensor) -> torch.Tensor: + """Convert GPU tensors -> contiguous CPU tensors if needed.""" + return maybe_gpu_tensor.detach().cpu().contiguous() + + def update_omni_tensor_prefix_cache( + self, + hidden_states: torch.Tensor | None, + multimodal_outputs: dict[str, torch.Tensor] | None, + num_tokens_unpadded: int, + slot_mapping: torch.Tensor, + num_tokens_padded: int | None = None, + ): + """Updates the hidden cache state for the provided hidden states and multimodal outputs. + + Args: + hidden_states: Hidden states tensor to cache (if any) + multimodal_outputs: Multimodal dict whose tensors may be cached + num_tokens_unpadded: Number of tokens without padding + slot_mapping: Slot mapping for the input sequence + num_tokens_padded: Total number of tokens including padding + """ + unpadded_slot_mapping = slot_mapping[:num_tokens_unpadded] + if num_tokens_padded is None: + num_tokens_padded = num_tokens_unpadded + + if hidden_states is not None: + # Slice to unpadded portion before caching + hidden_states = hidden_states[:num_tokens_unpadded] + # Ensure that hidden states are on the CPU + hidden_states = OmniTensorPrefixCache._coerce_to_cpu_tensor(hidden_states) + # View the cache as 2D so that we can treat our slots as row indices + flat_cache = self.hidden_states_cache.view(-1, self.hidden_states_cache.shape[-1]) + flat_cache[unpadded_slot_mapping] = hidden_states + logger.debug("Writing to hidden states for %s tokens", num_tokens_unpadded) + + # Do the same for the stage's cached multimodal outputs + if multimodal_outputs is not None: + # If we haven't initialized the keys already, do it now + # We check against the padded token count since we haven't sliced yet + self.maybe_init_missing_mm_cache_keys( + multimodal_outputs, + seq_len=num_tokens_padded, + ) + + for mm_out_key, mm_cache in self.mm_outputs_cache.items(): + if mm_out_key in multimodal_outputs: + # Slice to unpadded portion before caching + mm_state = multimodal_outputs[mm_out_key][:num_tokens_unpadded] + mm_state = OmniTensorPrefixCache._coerce_to_cpu_tensor(mm_state) + flat_cache = mm_cache.view(-1, mm_cache.shape[-1]) + flat_cache[unpadded_slot_mapping] = mm_state + logger.debug("Writing to mm output cache for %s tokens", num_tokens_unpadded) + + def _coerce_to_payload_dict( + self, + element: object, + query_start_loc: torch.Tensor, + input_batch: InputBatch, + num_scheduled_tokens: dict[str, int], + ) -> dict[str, object]: + """Build the multimodal passthrough data per request for + the object under consideration. This is identical to the case + for no prefix cache when we tensor does have a first dimension + matching the seq len. + """ + elem_dict = {} + for req_id in input_batch.req_ids: + req_idx = input_batch.req_id_to_index[req_id] + start = query_start_loc[req_idx] + end = start + num_scheduled_tokens[req_id] + elem_dict[req_id] = to_payload_element( + element, req_idx, start=start, end=end, pass_lists_through=True, seq_len=None + ) + return elem_dict + + def get_merged_multimodal_states( + self, + query_start_loc: torch.Tensor, + input_batch: InputBatch, + multimodal_outputs: dict, + num_scheduled_tokens: dict[str, int], + ): + """Get the merged multimodal states if hidden state prefix caching is enabled.""" + combined_multimodal_outputs = {} + # First get the prefix cached tensors that are present in the mm data + for mm_key in self.mm_cache_keys: + if mm_key in multimodal_outputs: + combined_multimodal_outputs[mm_key] = self._get_merged_tensors( + query_start_loc=query_start_loc, + input_batch=input_batch, + cache=self.mm_outputs_cache[mm_key], + hidden_states=multimodal_outputs[mm_key], + num_scheduled_tokens=num_scheduled_tokens, + ) + + # Then, get everything else (passthrough data); first, convert to CPU + # tensors similarly to the non prefix cached path, and then populate + # the subdicts mapping request IDs -> payload objects + passthrough_keys = set(multimodal_outputs.keys()) - self.mm_cache_keys + passthrough_mm_data = {k: v for k, v in multimodal_outputs.items() if k in passthrough_keys} + mm_cpu = build_mm_cpu(multimodal_outputs=passthrough_mm_data) + + for mm_key, mm_val in mm_cpu.items(): + combined_multimodal_outputs[mm_key] = self._coerce_to_payload_dict( + element=mm_val, + query_start_loc=query_start_loc, + input_batch=input_batch, + num_scheduled_tokens=num_scheduled_tokens, + ) + return combined_multimodal_outputs + + def get_merged_hidden_states(self, *args, **kwargs) -> dict[str, torch.Tensor]: + """Get the merged hidden states.""" + return self._get_merged_tensors( + *args, + **kwargs, + cache=self.hidden_states_cache, + ) + + def _get_merged_tensors( + self, + query_start_loc: torch.Tensor, + input_batch: InputBatch, + cache: torch.Tensor, + hidden_states: torch.Tensor, + num_scheduled_tokens: dict[str, int], + ) -> dict[str, torch.Tensor]: + """When hidden state caching is enabled, takes the input hidden_states, + which only correspond to the scheduled tokens, and returns a mapping + from request IDs to their full hidden states. This is accomplished by + looking up the block IDs & scheduled token counts to split the + hidden_states. + """ + # We do not support hybrid caches at the moment. + if len(input_batch.block_table.block_tables) > 1: + logger.warning_once( + "Omni prefix caching is enabled, but the batch block table appears to" + " have multiple kv groups; only the first group will be used!" + ) + + combined_hidden_states = {} + hidden_states = OmniTensorPrefixCache._coerce_to_cpu_tensor(hidden_states) + for req_id in input_batch.req_ids: + req_idx = input_batch.req_id_to_index[req_id] + + if req_id in self._new_req_cache_hit_ids: + block_ids = self._get_cached_block_ids(req_idx, input_batch) + cached_hs = cache[block_ids].reshape(-1, cache.shape[-1]) + + # Slice the hidden states corresponding to this request; + # we do this by using the query start + start = query_start_loc[req_idx] + new_hs = hidden_states[start : start + num_scheduled_tokens[req_id]] + combined_hidden_states[req_id] = torch.cat([cached_hs, new_hs], dim=0) + else: + # cache miss for this request, pass through normally + start = query_start_loc[req_idx] + new_hs = hidden_states[start : start + num_scheduled_tokens[req_id]] + combined_hidden_states[req_id] = new_hs + + return combined_hidden_states + + def _get_cached_block_ids(self, req_idx: int, input_batch: InputBatch) -> torch.Tensor: + """Given an input batch and request index in the batch (not ID), get the + block IDs corresponding to the cache hit. + """ + num_computed = input_batch.num_computed_tokens_cpu[req_idx] + # NOTE: vLLM only caches full blocks + num_cached_blocks = num_computed // self.block_size + # Get the block IDs attached to this cache hit and reindex into + # the flattened cached hidden states (i.e., 1 row per token). + return input_batch.block_table[0].block_table.cpu[req_idx, :num_cached_blocks] diff --git a/vllm_omni/utils/mm_outputs.py b/vllm_omni/utils/mm_outputs.py new file mode 100644 index 000000000..66d4e6ffe --- /dev/null +++ b/vllm_omni/utils/mm_outputs.py @@ -0,0 +1,93 @@ +"""Utilities for handling multimodal outputs / building multimodal output +payloads, most of which are shared by the prefix cache / no prefix cache path. +""" + +import torch +from vllm.logger import init_logger + +logger = init_logger(__name__) + + +def build_mm_cpu(multimodal_outputs: dict) -> dict[str, object]: + """Pre-copies multimodal tensor to CPU once (not per-request) to avoid + redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU. + + In the case of prefix caching, the multimodal outputs provided will + only contain the passthrough data. + + Args: + multimodal_outputs: Multimodal dict mapping strings to objects. + """ + # Pre-copy multimodal tensors to CPU once (not per-request) to avoid + # redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU. + mm_cpu: dict[str, object] = {} + # Currently there are some cases where this is true at the + # moment, which should be fixed. + if not isinstance(multimodal_outputs, dict): + logger.warning("Multimodal outputs are not a dict and will not be passed") + + if multimodal_outputs: + for k, v in multimodal_outputs.items(): + if isinstance(v, torch.Tensor): + mm_cpu[k] = v.detach().to("cpu").contiguous() + elif isinstance(v, dict): + sub_dict: dict[str, torch.Tensor] = {} + for sk, sv in v.items(): + if isinstance(sv, torch.Tensor): + sub_dict[str(sk)] = sv.detach().to("cpu").contiguous() + if sub_dict: + mm_cpu[k] = sub_dict + elif isinstance(v, list) and len(v) > 0: + cpu_list = [] + for elem in v: + if isinstance(elem, torch.Tensor): + cpu_list.append(elem.detach().to("cpu").contiguous()) + else: + cpu_list.append(elem) + mm_cpu[k] = cpu_list + elif v is not None: + mm_cpu[k] = v + return mm_cpu + + +def to_payload_element( + element: object, idx: int, start: int, end: int, pass_lists_through: bool = False, seq_len: int | None = None +): + """Build an mm payload element corresponding to one request index + from an element containing 0 or more CPU tensors. + + Args: + element: The object to be added to the payload. + idx: The index of the request. + start: The start index corresponding to the request idx. + end: The end index corresponding to the request idx. + pass_lists_through: bool Whether or not lists should be treated as + passthrough data; this should be False in normal cases, but True + if we need to avoid splitting nonempty lists prior to calling + postprocess, which is the case for prefix cache. + seq_len: Optional sequence length (i.e., dim 0 of hidden states). + This should be set to None in the prefix caching case, because + the condition that would be executed here is the same as the + criteria for being added to the multimodal outputs cache. + """ + # Prefix cache won't hit this case because this is the condition + # for being a mm_cache_key in the multimodal outputs tensor. + if seq_len is not None and isinstance(element, torch.Tensor) and element.shape[0] == seq_len: + return element[start:end].contiguous() + # Every other case is shared between prefix cache (passthrough data) + # and running a model without prefix caching. + elif isinstance(element, dict): + return {sk: sv[start:end].contiguous() for sk, sv in element.items()} + elif isinstance(element, list): + # For lists, clone tensors to avoid cross-request aliasing + if pass_lists_through: + return [elem.clone() if isinstance(elem, torch.Tensor) else elem for elem in element] + element = element[idx] if idx < len(element) else element[0] + if isinstance(element, torch.Tensor): + element = element.clone() + return element + elif isinstance(element, torch.Tensor): + # List-derived tensor payloads are request-invariant; clone to + # avoid accidental cross-request aliasing on downstream mutation. + return element.clone() + return element diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 62a0c8571..f37b2224e 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -39,6 +39,7 @@ from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm_omni.distributed.omni_connectors.kv_transfer_manager import OmniKVTransferManager from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.utils.mm_outputs import build_mm_cpu, to_payload_element from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner from vllm_omni.worker.omni_connector_model_runner_mixin import OmniConnectorModelRunnerMixin @@ -201,6 +202,63 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): finally: set_cudagraph_capturing_enabled(False) + def _maybe_update_prefix_cache( + self, + hidden_states: torch.Tensor, + multimodal_outputs: dict, + num_tokens_unpadded: int, + num_tokens_padded: int, + ): + """If prefix caching is enabled and it's the last pipeline parallelism rank, + retrieve the hidden states & multimodal outputs from the prefix cache based + on our batch slot mappings. + """ + # Cache hidden states if we've enabled hidden state prefix caching + # unless this isn't the last pipeline parallelism rank. + if self.omni_prefix_cache is not None and get_pp_group().is_last_rank: + # If this happens, it generally means the model is not following the correct + # interface yet and is therefore currently not compatible with prefix cache. + if multimodal_outputs is not None and not isinstance(multimodal_outputs, dict): + logger.warning_once( + "prefix caching expects mm outputs to be a dict, but got %s", + type(multimodal_outputs), + ) + + self.omni_prefix_cache.update_omni_tensor_prefix_cache( + hidden_states=hidden_states, + multimodal_outputs=multimodal_outputs, + num_tokens_unpadded=num_tokens_unpadded, + slot_mapping=self.input_batch.block_table[0].slot_mapping.cpu, + num_tokens_padded=num_tokens_padded, + ) + + def _maybe_get_combined_prefix_cache_tensors( + self, + hidden_states: torch.Tensor, + multimodal_outputs: dict, + num_scheduled_tokens: dict[str, int], + ) -> tuple[dict[str, torch.Tensor] | None, dict | None]: + """If prefix caching is enabled, extract the merged hidden states and multimodal outputs for + all requests in the batch (including those that aren't a hit on Prefix cache). + """ + # Prior to applying the post-processing func, extract + # the prefix cached hidden states and multimodal states. + combined_hidden_states, combined_multimodal_outputs = None, None + if self.omni_prefix_cache is not None: + combined_hidden_states = self.omni_prefix_cache.get_merged_hidden_states( + query_start_loc=self.query_start_loc.cpu, + input_batch=self.input_batch, + hidden_states=hidden_states, + num_scheduled_tokens=num_scheduled_tokens, + ) + combined_multimodal_outputs = self.omni_prefix_cache.get_merged_multimodal_states( + query_start_loc=self.query_start_loc.cpu, + input_batch=self.input_batch, + multimodal_outputs=multimodal_outputs, + num_scheduled_tokens=num_scheduled_tokens, + ) + return combined_hidden_states, combined_multimodal_outputs + @torch.inference_mode() def execute_model( self, @@ -476,6 +534,15 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): hidden_states, multimodal_outputs = self.extract_multimodal_outputs(model_output) + # Cache hidden states & multimodal outputs if we've enabled hidden state + # prefix caching unless this isn't the last pipeline parallelism rank. + self._maybe_update_prefix_cache( + hidden_states=hidden_states, + multimodal_outputs=multimodal_outputs, + num_tokens_unpadded=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded, + ) + if not self.broadcast_pp_output: # Common case. if not get_pp_group().is_last_rank: @@ -589,6 +656,23 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): return super()._sample(logits, spec_decode_metadata) + @staticmethod + def _resolve_req_hidden_states( + hidden_states_cpu: torch.Tensor, + combined_hidden_states: dict[str, torch.Tensor] | None, + rid: str, + start: int, + end: int, + ): + if combined_hidden_states is not None: + # We always have all request IDs for prefix cache, even for + # partial cache misses, so this should never happen. + if rid not in combined_hidden_states: + raise RuntimeError("Request IDs in the batch are missing from the merged states!") + return combined_hidden_states[rid] + # Prefix caching is disabled + return hidden_states_cpu[start:end] + @torch.inference_mode() def sample_tokens( self, @@ -597,6 +681,13 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): kv_extracted_req_ids = getattr(self, "kv_extracted_req_ids", None) self.kv_extracted_req_ids = None + # Used for prefix cache + combined_hidden_states = None + combined_multimodal_outputs = None + # Used when we don't use prefix cache; prefix cache builds the payloads + # internally since it already needs to do this for the cached tensors + mm_cpu = {} + if self.execute_model_state is None: kv_connector_output = self.kv_connector_output self.kv_connector_output = None @@ -628,6 +719,7 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): slot_mappings, # OMNI: unpack slot_mappings for drafter ) = self.execute_model_state self.execute_model_state = None + seq_len = hidden_states.shape[0] # Apply structured output bitmasks if present. if grammar_output is not None: @@ -749,37 +841,29 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): dtype=np.int32, ) - self._process_additional_information_updates( - hidden_states, multimodal_outputs, num_scheduled_tokens_np, scheduler_output - ) + # Prior to applying the post-processing func, extract + # the prefix cached hidden states and multimodal states. + if self.omni_prefix_cache is not None: + ( + combined_hidden_states, + combined_multimodal_outputs, + ) = self._maybe_get_combined_prefix_cache_tensors( + hidden_states, + multimodal_outputs, + scheduler_output.num_scheduled_tokens, + ) + # Otherwise we don't have the mm CPU data yet, so we still need to build it + if self.omni_prefix_cache is None: + mm_cpu = build_mm_cpu(multimodal_outputs) - # Pre-copy multimodal tensors to CPU once (not per-request) to avoid - # redundant D2H transfers when gpu_resident_buffer_keys keeps them on GPU. - mm_cpu: dict[str, object] = {} - if isinstance(multimodal_outputs, dict) and multimodal_outputs: - for k, v in multimodal_outputs.items(): - try: - if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: - mm_cpu[k] = v.detach().to("cpu").contiguous() - elif isinstance(v, dict): - sub_dict: dict[str, torch.Tensor] = {} - for sk, sv in v.items(): - if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]: - sub_dict[str(sk)] = sv.detach().to("cpu").contiguous() - if sub_dict: - mm_cpu[k] = sub_dict - elif isinstance(v, list): - if len(v) == 0: - continue - cpu_list = [] - for elem in v: - if isinstance(elem, torch.Tensor): - cpu_list.append(elem.detach().to("cpu").contiguous()) - else: - cpu_list.append(elem) - mm_cpu[k] = cpu_list - except Exception as e: - logger.error(f"Error in merge multimodal outputs: {e}") + self._process_additional_information_updates( + hidden_states, + multimodal_outputs, + num_scheduled_tokens_np, + scheduler_output, + combined_hidden_states, + combined_multimodal_outputs, + ) pooler_output: list[dict[str, object]] = [] for rid in req_ids_output_copy: @@ -787,29 +871,43 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): start = int(self.query_start_loc.cpu[idx]) sched = int(num_scheduled_tokens_np[idx]) end = start + sched - hidden_slice = hidden_states_cpu[start:end] - payload: dict[str, object] = {"hidden": hidden_slice} - if mm_cpu: - mm_payload: dict[str, object] = {} - for k, v in mm_cpu.items(): - if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: - mm_payload[k] = v[start:end].contiguous() - elif isinstance(v, dict): - mm_payload[k] = {sk: sv[start:end].contiguous() for sk, sv in v.items()} - elif isinstance(v, list): - element = v[idx] if idx < len(v) else v[0] - if element is not None: - if isinstance(element, torch.Tensor): - element = element.clone() - mm_payload[k] = element - # Skip None elements: msgspec cannot serialize None - # in dict[str, torch.Tensor] typed fields. - elif isinstance(v, torch.Tensor): - # List-derived tensor payloads are request-invariant; clone to - # avoid accidental cross-request aliasing on downstream mutation. - mm_payload[k] = v.clone() - else: - mm_payload[k] = v + # If prefix cache is enabled, we have already split everything + # by request and converted the states to CPU tensors + req_hidden_states = self._resolve_req_hidden_states( + hidden_states_cpu, + combined_hidden_states, + rid, + start, + end, + ) + payload: dict[str, object] = {"hidden": req_hidden_states} + + mm_payload: dict[str, object] = {} + if combined_multimodal_outputs or mm_cpu: + if combined_multimodal_outputs: + # Prefix cache enabled; all items have already been processed + # and split apart for each request as needed, and all tensors + # have already been detached to the CPU. The only exception is + # lists, which we keep as passthrough data for consistent behavior + # in postprocess. + for mm_key in combined_multimodal_outputs.keys(): + value = combined_multimodal_outputs[mm_key][rid] + if isinstance(value, list): + mm_payload[mm_key] = value[idx] if idx < len(value) else value[0] + else: + mm_payload[mm_key] = value + + else: + # Prefix cache disabled; we still need to process the data + for mm_key, mm_val in mm_cpu.items(): + mm_payload[mm_key] = to_payload_element( + element=mm_val, + idx=idx, + start=start, + end=end, + pass_lists_through=False, + seq_len=seq_len, + ) payload.update(mm_payload) pooler_output.append(payload) with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 5ff62c11b..de78011c7 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -20,6 +20,7 @@ from vllm.v1.worker.gpu_input_batch import CachedRequestState from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices +from vllm_omni.core.prefix_cache import OmniTensorPrefixCache from vllm_omni.engine.serialization import deserialize_additional_information from vllm_omni.model_executor.layers.rotary_embedding.mrope import OmniMRotaryEmbedding as MRotaryEmbedding from vllm_omni.model_executor.models.output_templates import OmniOutput @@ -43,6 +44,9 @@ class OmniGPUModelRunner(GPUModelRunner): self.model_intermediate_buffer: dict[str, dict[str, Any]] = {} self._omni_num_scheduled_tokens_np: np.ndarray | None = None self._omni_last_model_output: object | None = None + # The Omni tensor prefix cache will be allocated + # when we initialize the metadata builders if enabled + self.omni_prefix_cache = None def initialize_metadata_builders(self, kv_cache_config, kernel_block_sizes): """Override to fix scheduler_metadata buffer size for FA3 + CUDA graph. @@ -70,6 +74,16 @@ class OmniGPUModelRunner(GPUModelRunner): device=sm.device, ) + # Initialize the wrapper for both multimodal output tensors + # and for hidden states to be passed between stages + if self.cache_config.enable_prefix_caching: + self.omni_prefix_cache = OmniTensorPrefixCache( + num_blocks=kv_cache_config.num_blocks, + block_size=self.cache_config.block_size, + hidden_size=self.model_config.get_hidden_size(), + hs_dtype=self.dtype, + ) + @instrument(span_name="Loading (GPU)") def load_model(self, *args, **kwargs) -> None: super().load_model(*args, **kwargs) @@ -234,6 +248,10 @@ class OmniGPUModelRunner(GPUModelRunner): The SamplingMetadata is updated and copied to the GPU if there is a new/resumed/paused/finished request in the batch. """ + # Used for prefix cache + if self.omni_prefix_cache is not None: + self.omni_prefix_cache.reset_prefix_cached_new_req_ids() + # Remove finished requests from the cached states. for req_id in scheduler_output.finished_req_ids: self.requests.pop(req_id, None) @@ -294,6 +312,13 @@ class OmniGPUModelRunner(GPUModelRunner): reqs_to_add.append(req_state) continue + # Since this is the first time the request has been scheduled, + # num_computed_tokens > 0 means that we have a hit in prefix + # caching; mark it so that we can manage the hidden states + # later on as needed. + if self.omni_prefix_cache is not None and new_req_data.num_computed_tokens > 0: + self.omni_prefix_cache.add_prefix_cached_new_req_id(req_id) + sampling_params = new_req_data.sampling_params pooling_params = new_req_data.pooling_params @@ -1010,6 +1035,8 @@ class OmniGPUModelRunner(GPUModelRunner): multimodal_outputs: object, num_scheduled_tokens_np: np.ndarray, scheduler_output: "SchedulerOutput", + combined_hidden_states: dict[str, torch.Tensor] | None = None, + combined_multimodal_outputs: dict[str, object] | None = None, ) -> None: """Process model-provided per-request updates and merge into model_intermediate_buffer.""" try: @@ -1018,21 +1045,31 @@ class OmniGPUModelRunner(GPUModelRunner): if hasattr(self.model, "has_postprocess") and self.model.has_postprocess: for req_index, req_id in enumerate(self.input_batch.req_ids): req_infos = self.model_intermediate_buffer.get(req_id, {}) - start_offset = int(self.query_start_loc.cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - s, e = start_offset, start_offset + sched_tokens - # only consider to store data into update dict. - hidden_states_slice = hidden_states[s:e] + if combined_hidden_states: + # Combined hidden states contains all hidden states for every request + hidden_states_slice = combined_hidden_states[req_id] + else: + start_offset = int(self.query_start_loc.cpu[req_index]) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + s, e = start_offset, start_offset + sched_tokens + # only consider to store data into update dict. + hidden_states_slice = hidden_states[s:e] + + if combined_multimodal_outputs: + # NOTE this is a bit ugly, but the mm data is structured as a list of + # keys mapping to request IDs, and if enabled, we will always have all + # request IDs in every subdict, including for cache misses. + mm_out = {k: v[req_id] for k, v in combined_multimodal_outputs.items()} + else: + mm_out = multimodal_outputs update_dict = self.model.postprocess( - hidden_states_slice, multimodal_outputs=multimodal_outputs, **req_infos + hidden_states_slice, + multimodal_outputs=mm_out, + **req_infos, ) self._update_intermediate_buffer(req_id, update_dict) except Exception as e: - logger.error( - f"Error merging for requests:{self.input_batch.req_ids} " - f"additional information update: {e}, with the multimodal_outputs " - f"as {multimodal_outputs}" - ) + logger.error(f"Error merging for requests:{self.input_batch.req_ids} additional information update: {e}") import traceback traceback.print_exc() From e9581137e9d887c0876885d1c4a74ea7d63ba2eb Mon Sep 17 00:00:00 2001 From: Didan Deng <33117903+wtomin@users.noreply.github.com> Date: Thu, 16 Apr 2026 01:45:16 +0800 Subject: [PATCH 17/97] [Perf] Add Performance Test for Qwen-Image Step-Level Execution (#2707) Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> --- .../perf/tests/test_qwen_image_vllm_omni.json | 46 +++++++++++++++++++ 1 file changed, 46 insertions(+) diff --git a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json index 1f3a2bbf7..5ec7f1cc2 100644 --- a/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +++ b/tests/dfx/perf/tests/test_qwen_image_vllm_omni.json @@ -44,6 +44,52 @@ } ] }, + { + "test_name": "test_qwen_image_single_device_step_execution", + "description": "Single-device baseline (no parallelism) with step execution", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image", + "serve_args": { + "enable-diffusion-pipeline-profiler": true, + "step-execution": true + } + }, + "benchmark_params": [ + { + "name": "512x512_steps20", + "dataset": "random", + "task": "t2i", + "width": 512, + "height": 512, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.30, + "latency_mean": 3.50, + "peak_memory_mb_mean": 67000 + } + }, + { + "name": "1536x1536_steps35", + "dataset": "random", + "task": "t2i", + "width": 1536, + "height": 1536, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.037, + "latency_mean": 27.0, + "peak_memory_mb_mean": 74000 + } + } + ] + }, { "test_name": "test_qwen_image_ulysses2_cfg2_vae_patch4", "description": "Ulysses SP=2 + CFG-parallel=2 + VAE Patch Parallel=4", From 880a758b1f4b8be49618affbe4a735352f070993 Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Thu, 16 Apr 2026 10:15:39 +0800 Subject: [PATCH 18/97] [CI] Skip test_thinker_prefix_caching in tests/e2e/online_serving/test_qwen3_omni.py (#2836) Signed-off-by: wangyu <410167048@qq.com> --- tests/e2e/online_serving/test_qwen3_omni.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index c05f8f506..13af2ad11 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -183,6 +183,7 @@ def test_text_to_text_001(omni_server, openai_client) -> None: @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", prefix_test_params, indirect=True) +@pytest.mark.skip(reason="issue: #2833") def test_thinker_prefix_caching(omni_server, openai_client) -> None: """ Test thinker prefix caching by sending identical requests with an image (i.e., From c83f664fe17a372e0cfcf31b81b423ffee940e6b Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Thu, 16 Apr 2026 11:13:41 +0800 Subject: [PATCH 19/97] [CI][Perf] Add nightly PR labels, consolidate pipeline, and switch benchmark flag to --test-config-file (#2816) Signed-off-by: wangyu <410167048@qq.com> Co-authored-by: Y. Fisher Co-authored-by: inaniloquentee --- .buildkite/pipeline.yml | 12 +- .buildkite/test-nightly-diffusion.yml | 417 ---------------- .buildkite/test-nightly.yml | 446 ++++++++++++++++-- docs/contributing/ci/CI_5levels.md | 7 +- .../test_examples/l4_performance_tests.inc.md | 2 +- docs/contributing/ci/test_guide.md | 5 +- tests/dfx/conftest.py | 12 + tests/dfx/perf/scripts/run_benchmark.py | 49 +- .../perf/scripts/run_diffusion_benchmark.py | 25 +- .../tests/{test.json => test_qwen_omni.json} | 32 -- tests/dfx/perf/tests/test_tts.json | 34 ++ 11 files changed, 500 insertions(+), 541 deletions(-) delete mode 100644 .buildkite/test-nightly-diffusion.yml rename tests/dfx/perf/tests/{test.json => test_qwen_omni.json} (92%) create mode 100644 tests/dfx/perf/tests/test_tts.json diff --git a/.buildkite/pipeline.yml b/.buildkite/pipeline.yml index d9a231595..00823951d 100644 --- a/.buildkite/pipeline.yml +++ b/.buildkite/pipeline.yml @@ -44,11 +44,19 @@ steps: agents: queue: "cpu_queue_premerge" - # L4 Test — main+NIGHTLY=1 (scheduled), or PR with label nightly-test (e.g. add label then Rebuild) + # L4 Test — main+NIGHTLY=1 (scheduled), or PR with specific label (e.g. add label then Rebuild) - label: "Upload Nightly Pipeline" depends_on: image-build key: upload-nightly-pipeline - if: '(build.branch == "main" && build.env("NIGHTLY") == "1") || (build.branch != "main" && build.pull_request.labels includes "nightly-test")' + if: >- + (build.branch == "main" && build.env("NIGHTLY") == "1") || + (build.branch != "main" && ( + build.pull_request.labels includes "nightly-test" || + build.pull_request.labels includes "omni-test" || + build.pull_request.labels includes "tts-test" || + build.pull_request.labels includes "diffusion-x2iat-test" || + build.pull_request.labels includes "diffusion-x2v-test" + )) commands: - buildkite-agent pipeline upload .buildkite/test-nightly.yml agents: diff --git a/.buildkite/test-nightly-diffusion.yml b/.buildkite/test-nightly-diffusion.yml deleted file mode 100644 index b5ba8a117..000000000 --- a/.buildkite/test-nightly-diffusion.yml +++ /dev/null @@ -1,417 +0,0 @@ -# Nightly diffusion GPU tests — appended to the main nightly build via -# buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml -# from test-nightly.yml (step key: nightly-diffusion-model-test). Top-level groups are -# foldable in the Buildkite UI (Other / Wan / Qwen-Image). -env: - VLLM_WORKER_MULTIPROC_METHOD: spawn - HF_HUB_DOWNLOAD_TIMEOUT: 300 - HF_HUB_ETAG_TIMEOUT: 60 - -steps: - - group: ":card_index_dividers: Other Model Test" - key: nightly-other-model-test-group - steps: - - label: ":full_moon: Diffusion · Other · Function Test with H100" - timeout_in_minutes: 120 - # Shared nightly vs PR label conditional; referenced below as *nightly_or_pr_label - if: &nightly_or_pr_label 'build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"' - commands: - - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not test_qwen_image" -m "advanced_model and diffusion and H100" --run-level "advanced_model" - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - label: ":full_moon: Diffusion · Other · Function Test with L4" - timeout_in_minutes: 60 - if: *nightly_or_pr_label - commands: - - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and diffusion and L4" --run-level "advanced_model" - agents: - queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - shm-size: "8gb" - environment: - - "HF_HOME=/fsx/hf_cache" - - "HF_TOKEN" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" - - - label: ":full_moon: Diffusion · Other · Doc Test" - timeout_in_minutes: 60 - if: *nightly_or_pr_label - commands: - - export VLLM_TEST_CLEAN_GPU_MEMORY="1" - - pytest -s -v tests/examples/online_serving/test_text_to_image.py tests/examples/offline_inference/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model" - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - group: ":card_index_dividers: Wan Series Model Test" - key: nightly-wan-model-test-group - steps: - - label: ":full_moon: Diffusion · Wan · Function Test" - timeout_in_minutes: 90 - if: *nightly_or_pr_label - commands: - - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py -m "advanced_model" --run-level "advanced_model" - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - label: ":full_moon: Diffusion · Wan · Accuracy Test" - key: nightly-wan22-i2v-accuracy - timeout_in_minutes: 180 - if: *nightly_or_pr_label - commands: - - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - group: ":card_index_dividers: Qwen-Image Series Model Test" - key: nightly-qwen-image-edit-group - steps: - - label: ":full_moon: Diffusion · Qwen-Image · Function Test with H100" - timeout_in_minutes: 120 - if: *nightly_or_pr_label - commands: - - pytest -s -v tests/e2e/online_serving/test_qwen_image*_expansion.py -m "advanced_model and diffusion and H100" --run-level "advanced_model" - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 2 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - label: ":full_moon: Diffusion · Qwen-Image · GEBench Accuracy Test" - key: nightly-gebench-accuracy - timeout_in_minutes: 60 - if: *nightly_or_pr_label - commands: - - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1 - - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json" - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 1 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - label: ":full_moon: Diffusion · Qwen-Image · GEdit-Bench Accuracy Test" - key: nightly-gedit-bench-accuracy - timeout_in_minutes: 60 - if: *nightly_or_pr_label - commands: - - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1 - - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv" - - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json" - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 1 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE - value: "120" - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - label: ":full_moon: Diffusion · Qwen-Image · Perf Test" - key: nightly-qwen-image-performance - timeout_in_minutes: 180 - if: *nightly_or_pr_label - commands: - - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results - - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN - - export CACHE_DIT_VERSION=1.3.0 - # [HACK]: run upload in the same command block as pytest. - # Because `exit` aborts the entire commands list. - - | - set +e - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json - EXIT1=$$? - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json - EXIT2=$$? - pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json - EXIT3=$$? - if [ $$EXIT1 -eq 0 ] || [ $$EXIT2 -eq 0 ] || [ $$EXIT3 -eq 0 ]; then - buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" - fi - exit $$((EXIT1 | EXIT2 | EXIT3)) - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 4 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate - - - label: ":full_moon: Diffusion · Qwen-Image · Accuracy Test" - key: nightly-qwen-image-accuracy - timeout_in_minutes: 180 - if: *nightly_or_pr_label - commands: - - pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level advanced_model - agents: - queue: "mithril-h100-pool" - plugins: - - kubernetes: - podSpec: - containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - resources: - limits: - nvidia.com/gpu: 1 - volumeMounts: - - name: devshm - mountPath: /dev/shm - - name: hf-cache - mountPath: /root/.cache/huggingface - env: - - name: HF_HOME - value: /root/.cache/huggingface - - name: HF_TOKEN - valueFrom: - secretKeyRef: - name: hf-token-secret - key: token - nodeSelector: - node.kubernetes.io/instance-type: gpu-h100-sxm - volumes: - - name: devshm - emptyDir: - medium: Memory - - name: hf-cache - hostPath: - path: /mnt/hf-cache - type: DirectoryOrCreate diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 31b3e1797..58e1e55af 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -7,12 +7,11 @@ steps: # Group: collapses under one heading in the Buildkite UI; child steps still run in parallel. - group: ":card_index_dividers: Omni Model Test" key: nightly-omni-test-group + depends_on: upload-nightly-pipeline + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "omni-test" steps: - - label: ":full_moon: Omni · Function Test with H100" + - label: ":full_moon: Omni · Function Test" timeout_in_minutes: 90 - depends_on: upload-nightly-pipeline - # Shared nightly vs PR label conditional; referenced below as *nightly_or_pr_label - if: &nightly_or_pr_label 'build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test"' commands: - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and H100 and omni" --run-level "advanced_model" agents: @@ -49,31 +48,8 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - - label: ":full_moon: Omni · Function Test with L4" - timeout_in_minutes: 90 - depends_on: upload-nightly-pipeline - if: *nightly_or_pr_label - commands: - - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model" - agents: - queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU - plugins: - - docker#v5.2.0: - image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT - always-pull: true - propagate-environment: true - shm-size: "8gb" - environment: - - "HF_HOME=/fsx/hf_cache" - - "HF_TOKEN" - volumes: - - "/fsx/hf_cache:/fsx/hf_cache" - - label: ":full_moon: Omni · Doc Test with L4" timeout_in_minutes: 90 - depends_on: upload-nightly-pipeline - if: *nightly_or_pr_label commands: - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - pytest -s -v tests/examples/ -m "advanced_model and omni and L4" --run-level "advanced_model" @@ -93,8 +69,6 @@ steps: - label: ":full_moon: Omni · Doc Test with H100" timeout_in_minutes: 90 - depends_on: upload-nightly-pipeline - if: *nightly_or_pr_label commands: - pytest -s -v tests/examples/ -m "advanced_model and omni and H100" --run-level "advanced_model" agents: @@ -134,12 +108,9 @@ steps: - label: ":full_moon: Omni · Perf Test" key: nightly-omni-performance timeout_in_minutes: 180 - depends_on: upload-nightly-pipeline - if: *nightly_or_pr_label commands: - export BENCHMARK_DIR=tests/dfx/perf/results - - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py + - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_omni.json - buildkite-agent artifact upload "tests/dfx/perf/results/*.json" agents: queue: "mithril-h100-pool" @@ -175,23 +146,404 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate - # Dynamically appends steps from test-nightly-diffusion.yml into this build (same mechanism as - # pipeline.yml → test-ready.yml / test-merge.yml / test-nightly.yml). Foldable groups stay in the - # uploaded YAML (Other / Wan / Qwen-Image). - - label: ":card_index_dividers: Diffusion Model Test" - key: nightly-diffusion-model-test + + - group: ":card_index_dividers: TTS Model Test" + key: nightly-tts-test-group depends_on: upload-nightly-pipeline - if: *nightly_or_pr_label - commands: - - buildkite-agent pipeline upload .buildkite/test-nightly-diffusion.yml - agents: - queue: "cpu_queue_premerge" + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" || build.pull_request.labels includes "tts-test" + steps: + - label: ":full_moon: TTS · Function Test" + timeout_in_minutes: 90 + commands: + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -m "advanced_model and L4 and omni" --run-level "advanced_model" + agents: + queue: "gpu_1_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + + - label: ":full_moon: TTS · Perf Test" + key: nightly-tts-performance + timeout_in_minutes: 180 + commands: + - export BENCHMARK_DIR=tests/dfx/perf/results + - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" + - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json + - buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + # Diffusion X2I suite: x2i / x2a / x2t and related non-video paths; x2v is only in "Diffusion X2V Model Test" below. + - group: ":card_index_dividers: Diffusion X2I(&A&T) Model Test" + key: nightly-diffusion-x2iat-group + depends_on: upload-nightly-pipeline + if: >- + build.env("NIGHTLY") == "1" || + build.pull_request.labels includes "nightly-test" || + build.pull_request.labels includes "diffusion-x2iat-test" + steps: + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with H100" + timeout_in_minutes: 120 + commands: + - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and H100" --run-level "advanced_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: ":full_moon: Diffusion X2I(&A&T) · Function Test with L4" + timeout_in_minutes: 60 + commands: + - pytest -s -v tests/e2e/online_serving/test_*_expansion.py -k "not test_wan22_expansion and not test_wan_2_1_vace_expansion and not hunyuan" -m "advanced_model and diffusion and L4" --run-level "advanced_model" + agents: + queue: "gpu_4_queue" # g6.12xlarge instance on AWS, has 4 L4 GPU + plugins: + - docker#v5.2.0: + image: public.ecr.aws/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + always-pull: true + propagate-environment: true + shm-size: "8gb" + environment: + - "HF_HOME=/fsx/hf_cache" + - "HF_TOKEN" + volumes: + - "/fsx/hf_cache:/fsx/hf_cache" + + - label: ":full_moon: Diffusion X2I(&A&T) · Doc Test" + timeout_in_minutes: 60 + commands: + - export VLLM_TEST_CLEAN_GPU_MEMORY="1" + - pytest -s -v tests/examples/*/test_text_to_image.py -m "advanced_model and example and H100" --run-level "advanced_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: ":full_moon: Diffusion X2I(&A&T) · GEBench Accuracy Test" + timeout_in_minutes: 60 + commands: + - pytest -s -v tests/e2e/accuracy/test_gebench_h100_smoke.py --run-level advanced_model --gebench-model Qwen/Qwen-Image-2512 --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gebench-port 8093 --accuracy-workers 1 + - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gebench_qwen-image-2512/summary*.json" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: ":full_moon: Diffusion X2I(&A&T) · GEdit-Bench Accuracy Test" + timeout_in_minutes: 60 + commands: + - pytest -s -v tests/e2e/accuracy/test_gedit_bench_h100_smoke.py --run-level advanced_model --gedit-model Qwen/Qwen-Image-Edit --accuracy-judge-model QuantTrio/Qwen3-VL-30B-A3B-Instruct-AWQ --accuracy-gpu 0 --gedit-port 8093 --gedit-samples-per-group 20 --accuracy-workers 1 + - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_vie_score_*.csv" + - buildkite-agent artifact upload "tests/e2e/accuracy/artifacts/gedit_scores_qwen-image-edit/qwen-image-edit_all_all_summary_*.json" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE + value: "120" + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test" + key: nightly-diffusion-x2iat-performance + timeout_in_minutes: 180 + commands: + - export DIFFUSION_BENCHMARK_DIR=tests/dfx/perf/results + - export DIFFUSION_ATTENTION_BACKEND=FLASH_ATTN + - export CACHE_DIT_VERSION=1.3.0 + # [HACK]: run upload in the same command block as pytest. + # Because `exit` aborts the entire commands list. + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json + EXIT1=$$? + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_vllm_omni.json + EXIT2=$$? + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json + EXIT3=$$? + if [ $$EXIT1 -eq 0 ] || [ $$EXIT2 -eq 0 ] || [ $$EXIT3 -eq 0 ]; then + buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" + buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" + fi + exit $$((EXIT1 | EXIT2 | EXIT3)) + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 4 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + # Diffusion x2v only (Wan, HunyuanVideo, …). x2i/x2a/x2t live in the X2I group above, not here. + - group: ":card_index_dividers: Diffusion X2V Model Test" + key: nightly-diffusion-x2v-group + depends_on: upload-nightly-pipeline + if: >- + build.env("NIGHTLY") == "1" || + build.pull_request.labels includes "nightly-test" || + build.pull_request.labels includes "diffusion-x2v-test" + steps: + - label: ":full_moon: Diffusion X2V · Function Test" + timeout_in_minutes: 90 + commands: + - pytest -s -v tests/e2e/online_serving/test_wan22_expansion.py tests/e2e/online_serving/test_wan_2_1_vace_expansion.py tests/e2e/online_serving/test_hunyuan_video_15_expansion.py -m "advanced_model" --run-level "advanced_model" + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + + - label: ":full_moon: Diffusion X2V · Accuracy Test" + timeout_in_minutes: 180 + commands: + - pytest -s -v tests/e2e/accuracy/wan22_i2v/test_wan22_i2v_video_similarity.py --run-level advanced_model + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 2 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate - label: ":bar_chart: Testcase Statistics" key: nightly-testcase-statistics timeout_in_minutes: 120 depends_on: upload-nightly-pipeline - if: *nightly_or_pr_label + if: build.env("NIGHTLY") == "1" || build.pull_request.labels includes "nightly-test" commands: - python tools/nightly/buildkite_testcase_statistics.py -o tests/dfx/perf/results/buildkite_testcase_statistics.html - buildkite-agent artifact upload "tests/dfx/perf/results/*.html" @@ -234,15 +586,17 @@ steps: key: nightly-perf-distribution depends_on: - nightly-omni-performance - - nightly-qwen-image-performance + - nightly-tts-performance + - nightly-diffusion-x2iat-performance - nightly-testcase-statistics if: build.env("NIGHTLY") == "1" commands: - pip install openpyxl - export DEFAULT_INPUT_DIR=tests/dfx/perf/results - export DEFAULT_OUTPUT_DIR=tests/dfx/perf/results + - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-tts-performance - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-omni-performance - - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-qwen-image-performance + - buildkite-agent artifact download "tests/dfx/perf/results/*.json" . --step nightly-diffusion-x2iat-performance - buildkite-agent artifact download "tests/dfx/perf/results/*.html" . --step nightly-testcase-statistics - python tools/nightly/generate_nightly_perf_excel.py - python tools/nightly/generate_nightly_perf_html.py diff --git a/docs/contributing/ci/CI_5levels.md b/docs/contributing/ci/CI_5levels.md index 930603573..b0428ddd7 100644 --- a/docs/contributing/ci/CI_5levels.md +++ b/docs/contributing/ci/CI_5levels.md @@ -86,7 +86,8 @@ Through five levels (L1-L5) and common (Common) specifications, the system clari /tests/e2e/online_serving/test_{model_name}_expansion.py
/tests/e2e/offline_inference/test_{model_name}_expansion.py
Performance:
- /tests/dfx/perf/tests/test.json
+ /tests/dfx/perf/tests/test_qwen_omni.json (Omni), test_tts.json (TTS),
+ and /tests/dfx/perf/tests/test_{diffusion_model}_vllm_omni.json (Diffusion)
Doc Test:
tests/example/online_serving/test_{model_name}.py
tests/example/offline_inference/test_{model_name}.py @@ -530,13 +531,13 @@ L4 level testing is a comprehensive quality audit before a version release. It e ### 3.2 Testing Content and Scope - ***Full Functionality Testing***: Executes all test cases defined in `test_{model_name}_expansion.py`, covering all implemented features, positive flows, boundary conditions, and exception handling. -- ***Performance Testing***: Uses the `tests/dfx/perf/tests/test.json` configuration file to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization. +- ***Performance Testing***: Uses `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs in the form `tests/dfx/perf/tests/test_*_vllm_omni.json` (passed to `run_benchmark.py` via `--test-config-file`) to drive performance testing tools for stress, load, and endurance tests, collecting metrics like throughput, response time, and resource utilization. - ***Documentation Testing***: Verifies whether the example code provided to users is runnable and its results match the description. ### 3.3 Test Directory and Execution Files - ***Functional Testing***: Same directories as L3. -- ***Performance Test Configuration***: `tests/dfx/perf/tests/test.json` +- ***Performance Test Configuration***: `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, and diffusion configs `tests/dfx/perf/tests/test_*_vllm_omni.json` (e.g. `test_qwen_image_vllm_omni.json`) - ***Documentation Example Tests***: - - `tests/example/online_serving/test_{model_name}.py` - `tests/example/offline_inference/test_{model_name}.py` diff --git a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md index 8093e1459..f1f3073dc 100644 --- a/docs/contributing/ci/test_examples/l4_performance_tests.inc.md +++ b/docs/contributing/ci/test_examples/l4_performance_tests.inc.md @@ -1,4 +1,4 @@ -When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in tests/dfx/perf/tests/test.json: +When you want to add L4-level ***performance test*** cases, you can refer to the following format for case addition in `tests/dfx/perf/tests/test_qwen_omni.json`, `tests/dfx/perf/tests/test_tts.json`, or diffusion configs such as `tests/dfx/perf/tests/test_*_vllm_omni.json` (selected via `pytest ... run_benchmark.py --test-config-file `): ```JSON { diff --git a/docs/contributing/ci/test_guide.md b/docs/contributing/ci/test_guide.md index 425f24332..08b2e3b4e 100644 --- a/docs/contributing/ci/test_guide.md +++ b/docs/contributing/ci/test_guide.md @@ -45,7 +45,6 @@ Our test scripts use the pytest framework. First, please use `git clone https:// === "L3 level & L4 level" ```bash - cd tests pytest -s -v -m "advanced_model" --run-level=advanced_model ``` If you only want to run L3 test case, you can use: @@ -60,9 +59,9 @@ Our test scripts use the pytest framework. First, please use `git clone https:// ```bash pytest -s -v -m "core_model and distributed_cuda and L4" --run-level=core_model ``` - Note: To run performance tests, use: + Note: To run performance tests (defaults to ``test_qwen_omni.json``; use ``--test-config-file tests/dfx/perf/tests/test_tts.json`` for TTS): ```bash - pytest -s -v perf/scripts/run_benchmark.py + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py ``` The latest L3 test commands for various test suites can be found in the [pipeline](https://github.com/vllm-project/vllm-omni/blob/main/.buildkite/test-merge.yml). diff --git a/tests/dfx/conftest.py b/tests/dfx/conftest.py index e54141b34..997f25e6e 100644 --- a/tests/dfx/conftest.py +++ b/tests/dfx/conftest.py @@ -2,6 +2,8 @@ import json from pathlib import Path from typing import Any +import pytest + from tests.conftest import modify_stage_config @@ -95,3 +97,13 @@ def create_benchmark_indices( indices.append((test_name, idx)) return indices + + +def pytest_addoption(parser: pytest.Parser) -> None: + """Register shared CLI options for DFX benchmark suites.""" + parser.addoption( + "--test-config-file", + action="store", + default=None, + help=("Path to benchmark config JSON. Example: --test-config-file tests/dfx/perf/tests/test_tts.json"), + ) diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index 67dedcd04..d5ef1b49e 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -21,10 +21,30 @@ os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" -CONFIG_FILE_PATH = str(Path(__file__).parent.parent / "tests" / "test.json") -BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) -STAGE_INIT_TIMEOUT = 600 +def _get_config_file_from_argv() -> str | None: + """Read ``--test-config-file`` from ``sys.argv`` at import time so parametrization can use it.""" + import sys + for i, arg in enumerate(sys.argv): + if arg == "--test-config-file" and i + 1 < len(sys.argv): + return sys.argv[i + 1] + if arg.startswith("--test-config-file="): + return arg.split("=", 1)[1] + return None + + +_PERF_TESTS_DIR = Path(__file__).resolve().parent.parent / "tests" +_DEFAULT_CONFIG_FILE = str(_PERF_TESTS_DIR / "test_qwen_omni.json") + +CONFIG_FILE_PATH = _get_config_file_from_argv() +if CONFIG_FILE_PATH is None: + print( + "No --test-config-file in argv, using default: tests/dfx/perf/tests/test_qwen_omni.json " + "(override with e.g. --test-config-file tests/dfx/perf/tests/test_tts.json)" + ) + CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE + +BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) STAGE_CONFIGS_DIR = Path(__file__).parent.parent / "stage_configs" test_params = create_unique_server_params(BENCHMARK_CONFIGS, STAGE_CONFIGS_DIR) @@ -44,7 +64,7 @@ def omni_server(request): print(f"Starting OmniServer with test: {test_name}, model: {model}") - server_args = ["--stage-init-timeout", str(STAGE_INIT_TIMEOUT), "--init-timeout", "900"] + server_args = ["--stage-init-timeout", "300", "--init-timeout", "900"] if stage_config_path: server_args = ["--stage-configs-path", stage_config_path] + server_args with OmniServer(model, server_args) as server: @@ -97,8 +117,6 @@ def run_benchmark( ["vllm", "bench", "serve", "--omni"] + args + [ - "--num-warmups", - "2", "--save-result", "--result-dir", os.environ.get("BENCHMARK_DIR", "tests"), @@ -141,7 +159,6 @@ def run_benchmark( result["random_output_len"] = random_output_len with open(result_path, "w", encoding="utf-8") as f: json.dump(result, f, ensure_ascii=False, indent=2) - return result @@ -207,10 +224,6 @@ def _resolve_baseline_value( 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 @@ -245,14 +258,14 @@ def assert_result( ) -> None: assert result["completed"] == num_prompt, "Request failures exist" baseline_data = params.get("baseline", {}) - thresholds = _baseline_thresholds_for_step( - baseline_data, - sweep_index=sweep_index, - max_concurrency=max_concurrency, - request_rate=request_rate, - ) - for metric_name, baseline_value in thresholds.items(): + 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( diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 123f21405..8eeeec8df 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -5,8 +5,8 @@ Supports vLLM-Omni server backend: - vllm-omni (default): starts DiffusionServer via vllm_omni.entrypoints.cli.main, benchmarks with diffusion_benchmark_serving.py --backend vllm-omni -A config JSON file is REQUIRED via --config-file: - pytest run_diffusion_benchmark.py --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json +A config JSON file is REQUIRED via --test-config-file: + pytest run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json JSON config entries use a "server_type" field, and this runner executes the vllm-omni path. @@ -55,16 +55,16 @@ _BRANCHPOINT_COMMIT_SHA: str | None = None def _get_config_file_from_argv() -> str | None: - """Read --config-file from sys.argv at import time so pytest parametrize can use it. + """Read --test-config-file from sys.argv at import time so pytest parametrize can use it. pytest_addoption (below) registers the same flag so pytest does not reject it. - Supports both ``--config-file path`` and ``--config-file=path`` forms. + Supports both ``--test-config-file path`` and ``--test-config-file=path`` forms. Returns None if the flag is not present; callers must handle the missing case. """ for i, arg in enumerate(sys.argv): - if arg == "--config-file" and i + 1 < len(sys.argv): + if arg == "--test-config-file" and i + 1 < len(sys.argv): return sys.argv[i + 1] - if arg.startswith("--config-file="): + if arg.startswith("--test-config-file="): return arg.split("=", 1)[1] return None @@ -133,19 +133,6 @@ def _append_to_aggregated_file(record: dict[str, Any]) -> None: json.dump(records, f, indent=2, ensure_ascii=False) -# Register --config-file with pytest so it does not reject the argument. -def pytest_addoption(parser: pytest.Parser) -> None: - parser.addoption( - "--config-file", - action="store", - default=None, - help=( - "Path to the benchmark config JSON file (required). " - "Example: --config-file tests/dfx/perf/tests/test_qwen_image_vllm_omni.json" - ), - ) - - _server_lock = threading.Lock() # --------------------------------------------------------------------------- diff --git a/tests/dfx/perf/tests/test.json b/tests/dfx/perf/tests/test_qwen_omni.json similarity index 92% rename from tests/dfx/perf/tests/test.json rename to tests/dfx/perf/tests/test_qwen_omni.json index 159e27a06..4662f8c0c 100644 --- a/tests/dfx/perf/tests/test.json +++ b/tests/dfx/perf/tests/test_qwen_omni.json @@ -329,37 +329,5 @@ } } ] - }, - { - "test_name": "test_qwen3_tts", - "server_params": { - "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" - }, - "benchmark_params": [ - { - "dataset_name": "random", - "backend": "openai-audio-speech", - "endpoint": "/v1/audio/speech", - "num_prompts": [ - 10, - 40 - ], - "max_concurrency": [ - 1, - 4 - ], - "random_input_len": 100, - "random_output_len": 100, - "extra_body": { - "voice": "Vivian", - "language": "English" - }, - "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", - "baseline": { - "mean_audio_ttfp_ms": [6000, 6000], - "mean_audio_rtf": [0.3, 0.3] - } - } - ] } ] diff --git a/tests/dfx/perf/tests/test_tts.json b/tests/dfx/perf/tests/test_tts.json new file mode 100644 index 000000000..3583b45b4 --- /dev/null +++ b/tests/dfx/perf/tests/test_tts.json @@ -0,0 +1,34 @@ +[ + { + "test_name": "test_qwen3_tts", + "server_params": { + "model": "Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice" + }, + "benchmark_params": [ + { + "dataset_name": "random", + "backend": "openai-audio-speech", + "endpoint": "/v1/audio/speech", + "num_prompts": [ + 10, + 40 + ], + "max_concurrency": [ + 1, + 4 + ], + "random_input_len": 100, + "random_output_len": 100, + "extra_body": { + "voice": "Vivian", + "language": "English" + }, + "percentile-metrics": "ttft,e2el,audio_rtf,audio_ttfp,audio_duration", + "baseline": { + "mean_audio_ttfp_ms": [6000, 6000], + "mean_audio_rtf": [0.3, 0.3] + } + } + ] + } +] From de5f8a23b2cc4c51bdfe9d59f9887965c146d5d8 Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Thu, 16 Apr 2026 11:26:48 +0800 Subject: [PATCH 20/97] [Doc][Misc] Update DreamID-Omni Example; Add DreamID-Omni post process function (#2809) Signed-off-by: yuanheng --- .../offline_inference/x_to_video_audio.md | 28 ++++++++++++-- .../x_to_video_audio/x_to_video_audio.md | 28 ++++++++++++-- .../x_to_video_audio/x_to_video_audio.py | 38 +++++++++++++++---- .../dreamid_omni/pipeline_dreamid_omni.py | 15 ++++++++ vllm_omni/diffusion/registry.py | 1 + 5 files changed, 96 insertions(+), 14 deletions(-) diff --git a/docs/user_guide/examples/offline_inference/x_to_video_audio.md b/docs/user_guide/examples/offline_inference/x_to_video_audio.md index 8ea39d811..cec8d47c5 100644 --- a/docs/user_guide/examples/offline_inference/x_to_video_audio.md +++ b/docs/user_guide/examples/offline_inference/x_to_video_audio.md @@ -31,9 +31,9 @@ dreamid_omni/ ``` ### Run the Inference -``` +```python python x_to_video_audio.py \ - --model /xx/dreamid_omni \ + --model /path/to/dreamid_omni \ --prompt "Two people walking together and singing happily" \ --image-path ./example0.png ./example1.png \ --audio-path ./example0.wav ./example1.wav \ @@ -43,11 +43,33 @@ python x_to_video_audio.py \ --num-inference-steps 45 \ --height 704 \ --width 1280 \ - --output dreamid_omni.mp4 + --output out_dreamid_omni_twoip.mp4 ``` In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled. The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload. + +You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni + +For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing. + +```python +# Example usage for oneip, ref media from the official repo DreamID-Omni +python x_to_video_audio.py \ + --model /path/to/dreamid_omni \ + --prompt ": In the frame, a woman with black long hair is identified as .\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n maintains eye contact, swallows as if choosing her words, and says, I keep telling myself I’m fine,but some nights it feels like I’m just performing calm." \ + --image-path 9.png \ + --audio-path 9.wav \ + --video-negative-prompt "jitter, bad hands, blur, distortion" \ + --audio-negative-prompt "robotic, muffled, echo, distorted" \ + --cfg-parallel-size 2 \ + --num-inference-steps 45 \ + --height 704 \ + --width 1280 \ + --output out_dreamid_omni_oneip.mp4 +``` + + Key arguments: - `--prompt`: text description (string). - `--model`: path to the model local directory. diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md index 4b5188f41..13f2cfe7c 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.md +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.md @@ -30,9 +30,9 @@ dreamid_omni/ ``` ### Run the Inference -``` +```python python x_to_video_audio.py \ - --model /xx/dreamid_omni \ + --model /path/to/dreamid_omni \ --prompt "Two people walking together and singing happily" \ --image-path ./example0.png ./example1.png \ --audio-path ./example0.wav ./example1.wav \ @@ -42,11 +42,33 @@ python x_to_video_audio.py \ --num-inference-steps 45 \ --height 704 \ --width 1280 \ - --output dreamid_omni.mp4 + --output out_dreamid_omni_twoip.mp4 ``` In the current test scenario (2 images + 2 audio inputs), the VRAM requirement is 72GB, regardless of whether cfg-parallel is enabled or disabled. The VRAM usage can be reduced by enabling CPU offload via --enable-cpu-offload. + +You could take reference images/audios from the test cases in the official repo: https://github.com/Guoxu1233/DreamID-Omni + +For example, single IP ref resources can be found under https://github.com/Guoxu1233/DreamID-Omni/tree/main/test_case/oneip, you could download them correspondingly to your local and use them for testing. + +```python +# Example usage for oneip, ref media from the official repo DreamID-Omni +python x_to_video_audio.py \ + --model /path/to/dreamid_omni \ + --prompt ": In the frame, a woman with black long hair is identified as .\n**Overall Environment/Scene**: A lively open-kitchen café at night; stove flames flare, steam rises, and warm pendant lights swing slightly as staff move behind her. The shot is an upper-body close-up.\n**Main Characters/Subjects Appearance**: is a young woman with thick dark wavy hair and a side part. She wears a fitted black top under a light apron, a thin gold chain necklace, and small stud earrings.\n**Main Characters/Subjects Actions**: tastes the sauce with a spoon, then turns her face toward the camera while still holding the spoon, her expression shifting from focused to conflicted.\n maintains eye contact, swallows as if choosing her words, and says, I keep telling myself I’m fine,but some nights it feels like I’m just performing calm." \ + --image-path 9.png \ + --audio-path 9.wav \ + --video-negative-prompt "jitter, bad hands, blur, distortion" \ + --audio-negative-prompt "robotic, muffled, echo, distorted" \ + --cfg-parallel-size 2 \ + --num-inference-steps 45 \ + --height 704 \ + --width 1280 \ + --output out_dreamid_omni_oneip.mp4 +``` + + Key arguments: - `--prompt`: text description (string). - `--model`: path to the model local directory. diff --git a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py index 49a0f496f..322b184e5 100644 --- a/examples/offline_inference/x_to_video_audio/x_to_video_audio.py +++ b/examples/offline_inference/x_to_video_audio/x_to_video_audio.py @@ -5,10 +5,12 @@ import argparse import re import time +import numpy as np from PIL import Image from vllm.multimodal.media.audio import load_audio from vllm_omni.diffusion.data import DiffusionParallelConfig +from vllm_omni.diffusion.utils.media_utils import mux_video_audio_bytes from vllm_omni.entrypoints.omni import Omni from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -131,15 +133,35 @@ def main() -> None: if not outputs: raise RuntimeError("No output returned from DreamID-Omni.") - output = outputs[0].request_output - generated_video = output.images[0][0] - generated_audio = output.images[0][1] - try: - from dreamid_omni.utils.io_utils import save_video - except Exception as e: - raise RuntimeError(f"Failed to extract video and audio from DreamID-Omni output. Error: {e}") + result = outputs[0] + if not result.images: + raise RuntimeError("No video frames found in DreamID-Omni output.") + generated_video = result.images[0] + mm = result.multimodal_output or {} + generated_audio = mm.get("audio") + fps = int(mm.get("fps", 24)) + sample_rate = int(mm.get("audio_sample_rate", 16000)) + + # DreamID-Omni returns video as (C, F, H, W) float32 in [-1, 1]. + # mux_video_audio_bytes expects (F, H, W, C) uint8. + if not isinstance(generated_video, np.ndarray) or generated_video.ndim != 4: + raise RuntimeError(f"Unexpected video shape: {getattr(generated_video, 'shape', None)}") + frames = generated_video.transpose(1, 2, 3, 0) + frames = (np.clip((frames + 1.0) / 2.0, 0.0, 1.0) * 255.0).round().astype(np.uint8) + + audio_np = None + if generated_audio is not None: + audio_np = np.squeeze(np.asarray(generated_audio)).astype(np.float32) + output_path = args.output - save_video(output_path, generated_video, generated_audio, fps=24, sample_rate=16000) + video_bytes = mux_video_audio_bytes( + frames, + audio_np, + fps=float(fps), + audio_sample_rate=sample_rate, + ) + with open(output_path, "wb") as f: + f.write(video_bytes) print(f"Saved generated video to {output_path}") print(f"Total time: {elapsed:.2f}s") diff --git a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py index 974cc582f..c7ab4662d 100644 --- a/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py +++ b/vllm_omni/diffusion/models/dreamid_omni/pipeline_dreamid_omni.py @@ -38,6 +38,21 @@ from vllm_omni.diffusion.models.dreamid_omni.fusion import FusionModel logger = logging.getLogger(__name__) +def get_dreamid_omni_post_process_func(*args, **kwargs): + def post_process(output): + if isinstance(output, tuple) and len(output) == 2: + video, audio = output + return { + "video": video, + "audio": audio, + "audio_sample_rate": 16000, + "fps": 24, + } + return output + + return post_process + + AUDIO_CONFIG = { "patch_size": [1], "model_type": "t2a", diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 517b061ec..0bf8c0451 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -375,6 +375,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "HunyuanVideo15ImageToVideoPipeline": "get_hunyuan_video_15_i2v_post_process_func", "MagiHumanPipeline": "get_magi_human_post_process_func", "OmniVoicePipeline": "get_omnivoice_post_process_func", + "DreamIDOmniPipeline": "get_dreamid_omni_post_process_func", } _DIFFUSION_PRE_PROCESS_FUNCS = { From b43c6c6663311090e5a276826f2e2005d13ac05f Mon Sep 17 00:00:00 2001 From: Lancer Date: Thu, 16 Apr 2026 12:20:46 +0800 Subject: [PATCH 21/97] [Feat] add GLM-Image SP support (#1983) Signed-off-by: Lancer Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com> --- .../models/glm_image/test_glm_image_sp.py | 134 +++++++++ .../diffusion/attention/parallel/ulysses.py | 4 - .../models/glm_image/glm_image_transformer.py | 282 ++++++++++++++---- .../models/glm_image/pipeline_glm_image.py | 35 ++- 4 files changed, 394 insertions(+), 61 deletions(-) create mode 100644 tests/diffusion/models/glm_image/test_glm_image_sp.py diff --git a/tests/diffusion/models/glm_image/test_glm_image_sp.py b/tests/diffusion/models/glm_image/test_glm_image_sp.py new file mode 100644 index 000000000..1b1c8d7a7 --- /dev/null +++ b/tests/diffusion/models/glm_image/test_glm_image_sp.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Tests for GLM-Image Sequence Parallelism support.""" + +from unittest.mock import MagicMock, patch + +import pytest + +from vllm_omni.diffusion.data import DiffusionParallelConfig + + +@pytest.fixture(scope="function", autouse=True) +def setup_sp_groups(): + """Set up SP and TP groups for each test function.""" + with patch("vllm_omni.diffusion.distributed.parallel_state.get_sp_group") as mock_get_sp_group: + with patch("vllm.model_executor.layers.linear.get_tensor_model_parallel_world_size", return_value=1): + with patch("vllm.distributed.parallel_state.get_tp_group") as mock_get_tp_group: + mock_sp_group = MagicMock() + mock_sp_group.world_size = 4 + mock_get_sp_group.return_value = mock_sp_group + + mock_tp_group = MagicMock() + mock_tp_group.world_size = 1 + mock_get_tp_group.return_value = mock_tp_group + yield + + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_glm_image_sp_plan_defined(): + """Test that _sp_plan is properly defined on GlmImageTransformer2DModel.""" + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageTransformer2DModel, + ) + + assert hasattr(GlmImageTransformer2DModel, "_sp_plan") + plan = GlmImageTransformer2DModel._sp_plan + assert plan is not None + + # Verify plan structure + assert "prepare" in plan + assert "proj_out" in plan + + +def test_glm_image_sp_plan_valid(): + """Validate _sp_plan structure.""" + from vllm_omni.diffusion.distributed.sp_plan import validate_sp_plan + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageTransformer2DModel, + ) + + plan = GlmImageTransformer2DModel._sp_plan + validate_sp_plan(plan) + + +def test_glm_image_prepare_module_exists(): + """Test that GlmImagePrepare module exists.""" + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImagePrepare, + ) + + assert GlmImagePrepare is not None + + +def test_glm_image_attention_accepts_parallel_config(): + """Test that GlmImageAttention accepts parallel_config parameter.""" + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageAttention, + ) + + parallel_config = DiffusionParallelConfig( + ulysses_degree=2, + ring_degree=2, + tensor_parallel_size=1, + sequence_parallel_size=4, + ) + + attn = GlmImageAttention( + dim=2560, + num_heads=64, + head_dim=40, + parallel_config=parallel_config, + ) + + assert attn.parallel_config is not None + assert attn.parallel_config.sequence_parallel_size == 4 + + +def test_glm_image_transformer_block_accepts_parallel_config(): + """Test that GlmImageTransformerBlock accepts parallel_config parameter.""" + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageTransformerBlock, + ) + + parallel_config = DiffusionParallelConfig( + ulysses_degree=2, + ring_degree=2, + tensor_parallel_size=1, + sequence_parallel_size=4, + ) + + block = GlmImageTransformerBlock( + dim=2560, + num_attention_heads=64, + attention_head_dim=40, + time_embed_dim=512, + parallel_config=parallel_config, + ) + + assert block.attn1.parallel_config is not None + assert block.attn1.parallel_config.sequence_parallel_size == 4 + + +def test_glm_image_has_sp_support(): + """Test that GLM-Image has SP support implemented.""" + from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageTransformer2DModel, + ) + + # Check that the model has parallel_config support + assert hasattr(GlmImageTransformer2DModel, "__init__") + + # Verify the model can be instantiated with SP config + + # This test just verifies the structure exists + # Actual SP testing requires multi-GPU setup + + +@pytest.mark.cuda +@pytest.mark.sp +def test_glm_image_sp_inference(): + """Test SP inference (requires multi-GPU setup).""" + pytest.skip("Requires multi-GPU SP setup") diff --git a/vllm_omni/diffusion/attention/parallel/ulysses.py b/vllm_omni/diffusion/attention/parallel/ulysses.py index 5d860b335..326b5d456 100644 --- a/vllm_omni/diffusion/attention/parallel/ulysses.py +++ b/vllm_omni/diffusion/attention/parallel/ulysses.py @@ -414,10 +414,6 @@ class UlyssesParallelAttention: def post_attention(self, attn_output: torch.Tensor, ctx: ParallelAttentionContext | None) -> torch.Tensor: assert isinstance(ctx, _UlyssesCtx), f"Unexpected ctx type: {type(ctx)!r}" - # If we have joint tensors (Text), they were Head-Sliced. - # The main sequence (Image) was Sequence-Sliced. - # attn_output contains [Joint_Sliced | Image_Sliced] (if strategy='front'). - if ctx.joint_len > 0: joint_len = ctx.joint_len diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py index 490e0198b..7ff42a5f0 100644 --- a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py +++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py @@ -19,10 +19,16 @@ from vllm.model_executor.layers.linear import ( ) from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm_omni.diffusion.attention.backends.abstract import AttentionMetadata from vllm_omni.diffusion.attention.layer import Attention from vllm_omni.diffusion.cache.base import CachedTransformer -from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.distributed.hsdp_utils import is_transformer_block_module +from vllm_omni.diffusion.distributed.sp_plan import ( + SequenceParallelInput, + SequenceParallelOutput, +) +from vllm_omni.diffusion.forward_context import get_forward_context logger = init_logger(__name__) @@ -108,8 +114,8 @@ class GlmImageImageProjector(nn.Module): def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: batch_size, channel, height, width = hidden_states.shape - post_patch_height = height // self.patch_size - post_patch_width = width // self.patch_size + post_patch_height = torch.tensor(height // self.patch_size, device=hidden_states.device, dtype=torch.int64) + post_patch_width = torch.tensor(width // self.patch_size, device=hidden_states.device, dtype=torch.int64) # Reshape: [B, C, H, W] -> [B, H', W', C*p*p] -> [B, H'*W', C*p*p] hidden_states = hidden_states.reshape( @@ -159,6 +165,65 @@ class GlmImageRotaryPosEmbed(nn.Module): return (freqs.cos(), freqs.sin()) +class GlmImagePrepare(nn.Module): + """Prepare module for GLM-Image that handles patch embedding and RoPE computation. + + This module encapsulates the input processing pipeline to create a module boundary + where _sp_plan can shard outputs via split_output=True. + + Similar to Qwen-Image's ImageRopePrepare, this ensures hidden_states and RoPE + embeddings are sharded together to maintain dimension alignment. + """ + + def __init__( + self, + image_projector: nn.Module, + rope: GlmImageRotaryPosEmbed, + patch_size: int, + ): + super().__init__() + self.image_projector = image_projector + self.rope = rope + self.patch_size = patch_size + + def forward( + self, + hidden_states: torch.Tensor, + prior_hidden_states: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + """Process hidden_states and compute RoPE embeddings. + + Args: + hidden_states: Input latent tensor [B, C, H, W] + prior_hidden_states: Optional prior embedding to add + + Returns: + hidden_states: Patched hidden states [B, seq_len, D] + rope_cos: RoPE cos embeddings [seq_len, dim] + rope_sin: RoPE sin embeddings [seq_len, dim] + post_patch_height: Scalar tensor for height after patching + post_patch_width: Scalar tensor for width after patching + """ + batch_size, num_channels, height, width = hidden_states.shape + + post_patch_height = torch.tensor(height // self.patch_size, device=hidden_states.device, dtype=torch.int64) + post_patch_width = torch.tensor(width // self.patch_size, device=hidden_states.device, dtype=torch.int64) + + # Compute RoPE (uses original 4D hidden_states shape) + image_rotary_emb = self.rope(hidden_states) + rope_cos = image_rotary_emb[0].to(hidden_states.device) + rope_sin = image_rotary_emb[1].to(hidden_states.device) + + # Patch embedding: [B, C, H, W] -> [B, seq_len, D] + hidden_states = self.image_projector(hidden_states) + + # Add prior embedding if provided + if prior_hidden_states is not None: + hidden_states = hidden_states + prior_hidden_states + + return hidden_states, rope_cos, rope_sin, post_patch_height, post_patch_width + + class GlmImageAdaLayerNormZero(nn.Module): """Adaptive LayerNorm with zero initialization for both image and text streams.""" @@ -397,6 +462,7 @@ class GlmImageAttention(nn.Module): dim: int, num_heads: int, head_dim: int, + parallel_config: DiffusionParallelConfig | None = None, out_bias: bool = True, eps: float = 1e-5, ): @@ -404,6 +470,7 @@ class GlmImageAttention(nn.Module): self.dim = dim self.total_num_heads = num_heads self.head_dim = head_dim + self.parallel_config = parallel_config # QKV projection (fused for efficiency) self.to_qkv = QKVParallelLinear( @@ -450,16 +517,19 @@ class GlmImageAttention(nn.Module): attention_mask: torch.Tensor | None = None, kv_cache: GlmImageLayerKVCache | None = None, kv_cache_mode: KVCacheMode | None = None, + hidden_states_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Forward pass for joint attention. Args: - hidden_states: Image hidden states [B, img_seq_len, D] - encoder_hidden_states: Text hidden states [B, text_seq_len, D] - image_rotary_emb: Tuple of (cos, sin) for RoPE + hidden_states: Image hidden states [B, img_seq_len, D] (sharded in SP mode) + encoder_hidden_states: Text hidden states [B, text_seq_len, D] (full in SP mode) + image_rotary_emb: Tuple of (cos, sin) for RoPE (sharded in SP mode) + attention_mask: Optional attention mask kv_cache: Optional layer KV cache for image editing kv_cache_mode: Cache mode (WRITE, READ, SKIP) + hidden_states_mask: Mask for SP padding (True=valid, False=padding) Returns: Tuple of (image_hidden_states, text_hidden_states) @@ -467,6 +537,13 @@ class GlmImageAttention(nn.Module): dtype = encoder_hidden_states.dtype batch_size, text_seq_length, _ = encoder_hidden_states.shape + # Check if SP is enabled + sp_size = self.parallel_config.sequence_parallel_size if self.parallel_config else None + use_sp = sp_size is not None and sp_size > 1 + if use_sp: + forward_ctx = get_forward_context() + use_sp = not forward_ctx.split_text_embed_in_sp + # Concatenate text and image: [text, image] hidden_states_combined = torch.cat([encoder_hidden_states, hidden_states], dim=1) @@ -485,41 +562,88 @@ class GlmImageAttention(nn.Module): query = self.norm_q(query).to(dtype=dtype) key = self.norm_k(key).to(dtype=dtype) - # Apply RoPE only to image tokens (not text tokens) - if image_rotary_emb is not None: - # Only apply RoPE to image part (after text_seq_length) - query_img = query[:, text_seq_length:, :, :] - key_img = key[:, text_seq_length:, :, :] - from diffusers.models.embeddings import apply_rotary_emb + if use_sp: + # SP mode: use joint attention mechanism + # Split Q/K/V into text and image parts + text_query = query[:, :text_seq_length, :, :] + text_key = key[:, :text_seq_length, :, :] + text_value = value[:, :text_seq_length, :, :] + img_query = query[:, text_seq_length:, :, :] + img_key = key[:, text_seq_length:, :, :] + img_value = value[:, text_seq_length:, :, :] - query_img = apply_rotary_emb(query_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) - key_img = apply_rotary_emb(key_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) - query = torch.cat([query[:, :text_seq_length, :, :], query_img], dim=1) - key = torch.cat([key[:, :text_seq_length, :, :], key_img], dim=1) + # Apply RoPE only to image part + if image_rotary_emb is not None: + from diffusers.models.embeddings import apply_rotary_emb - # Handle KV cache for image editing - if kv_cache is not None and kv_cache_mode is not None: - if kv_cache_mode == KVCacheMode.WRITE: - kv_cache.store(key, value) - elif kv_cache_mode == KVCacheMode.READ: - k_cached, v_cached = kv_cache.get() - if k_cached is not None: - key = torch.cat([k_cached, key], dim=1) - value = torch.cat([v_cached, value], dim=1) - # KVCacheMode.SKIP: do nothing + img_query = apply_rotary_emb(img_query, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) + img_key = apply_rotary_emb(img_key, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) - # Attention computation - hidden_states_out = self.attn(query, key, value) - hidden_states_out = hidden_states_out.flatten(2, 3) - hidden_states_out = hidden_states_out.to(dtype) + # Create attention metadata for joint attention + attn_metadata = AttentionMetadata( + joint_query=text_query, + joint_key=text_key, + joint_value=text_value, + joint_strategy="front", + ) - # Output projection - for module in self.to_out: - hidden_states_out = module(hidden_states_out) + # Add padding mask for SP if available + if hidden_states_mask is not None: + attn_metadata.attn_mask = hidden_states_mask - # Split back to text and image - encoder_hidden_states_out = hidden_states_out[:, :text_seq_length, :] - hidden_states_out = hidden_states_out[:, text_seq_length:, :] + # Attention computation with joint text/image + # Note: Ulysses post_attention returns [text, image] concatenated + joint_hidden_states_out = self.attn(img_query, img_key, img_value, attn_metadata) + + # Project combined [text, image] outputs, then split. + # This keeps SP numerically aligned with the non-SP path. + joint_hidden_states_out = joint_hidden_states_out.flatten(2, 3).to(dtype) + for module in self.to_out: + joint_hidden_states_out = module(joint_hidden_states_out) + + encoder_hidden_states_out = joint_hidden_states_out[:, :text_seq_length, :] + hidden_states_out = joint_hidden_states_out[:, text_seq_length:, :] + else: + # Non-SP mode: original logic + # Apply RoPE only to image tokens (not text tokens) + if image_rotary_emb is not None: + query_img = query[:, text_seq_length:, :, :] + key_img = key[:, text_seq_length:, :, :] + from diffusers.models.embeddings import apply_rotary_emb + + query_img = apply_rotary_emb(query_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) + key_img = apply_rotary_emb(key_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) + query = torch.cat([query[:, :text_seq_length, :, :], query_img], dim=1) + key = torch.cat([key[:, :text_seq_length, :, :], key_img], dim=1) + + # Handle KV cache for image editing + if kv_cache is not None and kv_cache_mode is not None: + if kv_cache_mode == KVCacheMode.WRITE: + kv_cache.store(key, value) + elif kv_cache_mode == KVCacheMode.READ: + k_cached, v_cached = kv_cache.get() + if k_cached is not None: + key = torch.cat([k_cached, key], dim=1) + value = torch.cat([v_cached, value], dim=1) + + # Attention computation + attn_metadata = None + if attention_mask is not None: + if attention_mask.dim() == 3: + attention_mask = attention_mask.unsqueeze(1) + attn_metadata = AttentionMetadata(attn_mask=attention_mask) + + hidden_states_out = self.attn(query, key, value, attn_metadata) + hidden_states_out = hidden_states_out.flatten(2, 3) + hidden_states_out = hidden_states_out.to(dtype) + + # Output projection + for module in self.to_out: + hidden_states_out = module(hidden_states_out) + + # Split back to text and image + encoder_hidden_states_out = hidden_states_out[:, :text_seq_length, :] + hidden_states_out = hidden_states_out[:, text_seq_length:, :] return hidden_states_out, encoder_hidden_states_out @@ -628,6 +752,7 @@ class GlmImageTransformerBlock(nn.Module): attention_head_dim: int = 40, time_embed_dim: int = 512, ffn_hidden_dim: int | None = None, + parallel_config: DiffusionParallelConfig | None = None, ) -> None: super().__init__() @@ -637,6 +762,7 @@ class GlmImageTransformerBlock(nn.Module): dim=dim, num_heads=num_attention_heads, head_dim=attention_head_dim, + parallel_config=parallel_config, ) # 2. Feedforward @@ -654,6 +780,7 @@ class GlmImageTransformerBlock(nn.Module): attention_kwargs: dict[str, Any] | None = None, kv_cache: GlmImageLayerKVCache | None = None, kv_cache_mode: KVCacheMode | None = None, + hidden_states_mask: torch.Tensor | None = None, ) -> tuple[torch.Tensor, torch.Tensor]: """ Forward pass for transformer block. @@ -667,6 +794,7 @@ class GlmImageTransformerBlock(nn.Module): attention_kwargs: Additional attention arguments kv_cache: Layer-specific KV cache for image editing kv_cache_mode: Cache mode (WRITE, READ, SKIP) + hidden_states_mask: Mask for SP padding (True=valid, False=padding) Returns: Tuple of (image_hidden_states, text_hidden_states) @@ -693,6 +821,7 @@ class GlmImageTransformerBlock(nn.Module): attention_mask=attention_mask, kv_cache=kv_cache, kv_cache_mode=kv_cache_mode, + hidden_states_mask=hidden_states_mask, ) hidden_states = hidden_states + attn_hidden_states * gate_msa.unsqueeze(1) encoder_hidden_states = encoder_hidden_states + attn_encoder_hidden_states * c_gate_msa.unsqueeze(1) @@ -724,6 +853,26 @@ class GlmImageTransformer2DModel(CachedTransformer): """ _repeated_blocks = ["GlmImageTransformerBlock"] + # SP plan using GlmImagePrepare module for sharding hidden_states and RoPE together. + # Similar to Qwen-Image's ImageRopePrepare, this creates a module boundary where + # _sp_plan can shard outputs via split_output=True. + # + # Key insight: hidden_states and RoPE embeddings MUST be sharded together + # to maintain dimension alignment for RoPE computation in attention layers. + _sp_plan = { + # Shard GlmImagePrepare outputs (hidden_states and RoPE must be sharded together) + "prepare": { + # hidden_states: [B, seq_len, D] - shard along sequence dimension + 0: SequenceParallelInput(split_dim=1, expected_dims=3, split_output=True, auto_pad=True), + # RoPE cos: [seq_len, dim] - shard along sequence dimension + 1: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), + # RoPE sin: [seq_len, dim] - shard along sequence dimension + 2: SequenceParallelInput(split_dim=0, expected_dims=2, split_output=True, auto_pad=True), + # post_patch_height and post_patch_width are scalars, not sharded + }, + # Gather output at proj_out + "proj_out": SequenceParallelOutput(gather_dim=1, expected_dims=3), + } _hsdp_shard_conditions = [is_transformer_block_module] @@ -790,6 +939,9 @@ class GlmImageTransformer2DModel(CachedTransformer): dim=inner_dim, dim_out=inner_dim, inner_dim=inner_dim, activation_fn="linear-silu" ) + # Prepare module for SP (encapsulates patch embedding and RoPE for _sp_plan) + self.prepare = GlmImagePrepare(self.image_projector, self.rope, patch_size) + self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings( embedding_dim=time_embed_dim, condition_dim=condition_dim, @@ -806,6 +958,7 @@ class GlmImageTransformer2DModel(CachedTransformer): attention_head_dim, time_embed_dim, ffn_hidden_dim=ffn_hidden_dim, + parallel_config=self.parallel_config, ) for _ in range(num_layers) ] @@ -859,33 +1012,51 @@ class GlmImageTransformer2DModel(CachedTransformer): # Get KV cache mode kv_cache_mode = kv_cache.mode if kv_cache is not None else None - # 1. RoPE - if image_rotary_emb is None: - image_rotary_emb = self.rope(hidden_states) - # Move to correct device - image_rotary_emb = ( - image_rotary_emb[0].to(hidden_states.device), - image_rotary_emb[1].to(hidden_states.device), - ) + # Set SP context if enabled + sp_size = self.parallel_config.sequence_parallel_size + if sp_size is not None and sp_size > 1: + get_forward_context().split_text_embed_in_sp = False - # 2. Patch & Timestep embeddings - p = self.patch_size - post_patch_height = height // p - post_patch_width = width // p - - hidden_states = self.image_projector(hidden_states) + # Text embedding projection encoder_hidden_states = self.glyph_projector(encoder_hidden_states) # Prior embedding with dropout prior_embedding = self.prior_token_embedding(prior_token_id) prior_embedding[prior_token_drop] *= 0.0 prior_hidden_states = self.prior_projector(prior_embedding) - hidden_states = hidden_states + prior_hidden_states + + # 1. Prepare hidden_states and RoPE via GlmImagePrepare module + # _sp_plan will shard hidden_states and RoPE together via split_output=True + hidden_states, rope_cos, rope_sin, post_patch_height_t, post_patch_width_t = self.prepare( + hidden_states, prior_hidden_states + ) + image_rotary_emb = (rope_cos, rope_sin) + post_patch_height = int(post_patch_height_t.item()) + post_patch_width = int(post_patch_width_t.item()) # Timestep conditioning temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype) - # 3. Transformer blocks + # Create padding mask for SP if needed (after _sp_plan hooks have run) + hidden_states_mask = None + if sp_size is not None and sp_size > 1: + from vllm_omni.diffusion.forward_context import is_forward_context_available + + if is_forward_context_available(): + ctx = get_forward_context() + if ctx.sp_original_seq_len is not None and ctx.sp_padding_size > 0: + img_padded_seq_len = ctx.sp_original_seq_len + ctx.sp_padding_size + hidden_states_mask = torch.ones( + batch_size, + img_padded_seq_len, + dtype=torch.bool, + device=hidden_states.device, + ) + hidden_states_mask[:, ctx.sp_original_seq_len :] = False + if hidden_states_mask.all(): + hidden_states_mask = None + + # 2. Transformer blocks for layer_idx, block in enumerate(self.transformer_blocks): # Get layer-specific KV cache if available layer_kv_cache = kv_cache[layer_idx] if kv_cache is not None else None @@ -899,13 +1070,16 @@ class GlmImageTransformer2DModel(CachedTransformer): attention_kwargs, kv_cache=layer_kv_cache, kv_cache_mode=kv_cache_mode, + hidden_states_mask=hidden_states_mask, ) - # 4. Output norm & projection + # 3. Output norm & projection + # _sp_plan will gather hidden_states via proj_out hook hidden_states = self.norm_out(hidden_states, temb) hidden_states = self.proj_out(hidden_states) - # 5. Unpatchify: [B, H'*W', C*p*p] -> [B, C, H, W] + # 4. Unpatchify: [B, H'*W', C*p*p] -> [B, C, H, W] + p = self.patch_size hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, -1, p, p) output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py index 375f7e7b8..038636499 100644 --- a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py +++ b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py @@ -712,6 +712,14 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): if img is not None: preprocessed_images = [img] + # Priority: prompt dict (from ar2diffusion) > sampling_params + # ar2diffusion returns adjusted height/width that matches prior_token_ids + if not isinstance(first_prompt, str): + ar_height = first_prompt.get("height") + ar_width = first_prompt.get("width") + else: + ar_height = ar_width = None + img_height = req.sampling_params.height img_width = req.sampling_params.width @@ -719,12 +727,19 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): # Treat that as t2i warmup to avoid requiring i2i-only KV-cache inputs. is_image_edit = (preprocessed_images is not None) and (not is_dummy_warmup) - # Use image dimensions as default if available - height = req.sampling_params.height or img_height or self.default_sample_size * self.vae_scale_factor - width = req.sampling_params.width or img_width or self.default_sample_size * self.vae_scale_factor + # Use prompt dict dimensions (from ar2diffusion) as priority, then sampling_params + height = ( + ar_height or req.sampling_params.height or img_height or self.default_sample_size * self.vae_scale_factor + ) + width = ar_width or req.sampling_params.width or img_width or self.default_sample_size * self.vae_scale_factor num_inference_steps = req.sampling_params.num_inference_steps or 50 guidance_scale = req.sampling_params.guidance_scale or 1.5 + # Ensure dimensions are multiples of vae_scale_factor * patch_size + multiple_of = self.vae_scale_factor * self._patch_size + height = height // multiple_of * multiple_of + width = width // multiple_of * multiple_of + self.check_inputs(prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds) batch_size = 1 @@ -753,6 +768,20 @@ class GlmImagePipeline(nn.Module, DiffusionPipelineProfilerMixin): prior_token_id = prior_token_id.to(device=self.device, dtype=torch.long) if prior_token_id.dim() == 1: prior_token_id = prior_token_id.unsqueeze(0) + + # Validate that prior_token_id seq_len matches dimensions + prior_seq_len = prior_token_id.shape[1] + expected_seq_len = (height // self.vae_scale_factor // self._patch_size) * ( + width // self.vae_scale_factor // self._patch_size + ) + if prior_seq_len != expected_seq_len: + raise ValueError( + f"prior_token_ids seq_len ({prior_seq_len}) doesn't match dimensions " + f"({height}x{width}, expected seq_len={expected_seq_len}). " + f"This indicates a mismatch between AR output and Diffusion input. " + f"Please ensure ar2diffusion returns correct height/width." + ) + prior_token_image_ids = None if external_prior_image_ids is not None: if isinstance(external_prior_image_ids, torch.Tensor): From 24e61f4d7bccb61d020f9020c22da51546a4c7c5 Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Thu, 16 Apr 2026 14:52:53 +0800 Subject: [PATCH 22/97] [CI] add qwen image and layered accuracy test (#2772) Signed-off-by: david6666666 <530634352@qq.com> --- tests/e2e/accuracy/test_qwen_image.py | 124 ++++++++++++++ tests/e2e/accuracy/test_qwen_image_layered.py | 151 ++++++++++++++++++ tests/e2e/accuracy/utils.py | 47 ++++-- 3 files changed, 313 insertions(+), 9 deletions(-) create mode 100644 tests/e2e/accuracy/test_qwen_image.py create mode 100644 tests/e2e/accuracy/test_qwen_image_layered.py diff --git a/tests/e2e/accuracy/test_qwen_image.py b/tests/e2e/accuracy/test_qwen_image.py new file mode 100644 index 000000000..e73195017 --- /dev/null +++ b/tests/e2e/accuracy/test_qwen_image.py @@ -0,0 +1,124 @@ +from __future__ import annotations + +import base64 +import gc +import io +import os +from pathlib import Path + +import pytest +import requests +import torch +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from PIL import Image + +from tests.conftest import ( + OmniServer, + _run_post_test_cleanup, + _run_pre_test_cleanup, +) +from tests.e2e.accuracy.utils import assert_similarity, model_output_dir +from tests.utils import hardware_test + +MODEL_ID = "Qwen/Qwen-Image" +MODEL_ENV_VAR = "QWEN_IMAGE_MODEL" +PROMPT = "A photo of a cat sitting on a laptop keyboard, digital art style." +NEGATIVE_PROMPT = "blurry, low quality" +WIDTH = 512 +HEIGHT = 512 +NUM_INFERENCE_STEPS = 20 +TRUE_CFG_SCALE = 4.0 +SEED = 42 +SSIM_THRESHOLD = 0.97 +PSNR_THRESHOLD = 30.0 + + +def _model_name() -> str: + return os.environ.get(MODEL_ENV_VAR, MODEL_ID) + + +def _local_files_only(model: str) -> bool: + return Path(model).exists() + + +def _run_vllm_omni_qwen_image(*, model: str, output_path: Path) -> Image.Image: + server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"] + with OmniServer(model, server_args, use_omni=True) as omni_server: + response = requests.post( + f"http://{omni_server.host}:{omni_server.port}/v1/images/generations", + json={ + "model": omni_server.model, + "prompt": PROMPT, + "size": f"{WIDTH}x{HEIGHT}", + "n": 1, + "response_format": "b64_json", + "negative_prompt": NEGATIVE_PROMPT, + "num_inference_steps": NUM_INFERENCE_STEPS, + "true_cfg_scale": TRUE_CFG_SCALE, + "seed": SEED, + }, + timeout=600, + ) + response.raise_for_status() + payload = response.json() + assert len(payload["data"]) == 1 + image_bytes = base64.b64decode(payload["data"][0]["b64_json"]) + image = Image.open(io.BytesIO(image_bytes)).convert("RGB") + image.load() + image.save(output_path) + return image + + +def _run_diffusers_qwen_image(*, model: str, output_path: Path) -> Image.Image: + _run_pre_test_cleanup(enable_force=True) + pipe: DiffusionPipeline | None = None + try: + pipe = DiffusionPipeline.from_pretrained( + model, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + local_files_only=_local_files_only(model), + ).to("cuda") + generator = torch.Generator(device="cuda").manual_seed(SEED) + result = pipe( # pyright: ignore[reportCallIssue] + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + width=WIDTH, + height=HEIGHT, + num_inference_steps=NUM_INFERENCE_STEPS, + true_cfg_scale=TRUE_CFG_SCALE, + generator=generator, + ) + output_image = result.images[0].convert("RGB") + output_image.save(output_path) + return output_image + finally: + if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"): + pipe.maybe_free_model_hooks() + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + _run_post_test_cleanup(enable_force=True) + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_qwen_image_matches_diffusers(accuracy_artifact_root: Path) -> None: + model = _model_name() + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + + vllm_output = _run_vllm_omni_qwen_image(model=model, output_path=output_dir / "vllm_omni.png") + diffusers_output = _run_diffusers_qwen_image(model=model, output_path=output_dir / "diffusers.png") + + assert_similarity( + model_name=MODEL_ID, + vllm_image=vllm_output, + diffusers_image=diffusers_output, + width=WIDTH, + height=HEIGHT, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + ) diff --git a/tests/e2e/accuracy/test_qwen_image_layered.py b/tests/e2e/accuracy/test_qwen_image_layered.py new file mode 100644 index 000000000..04b13df3b --- /dev/null +++ b/tests/e2e/accuracy/test_qwen_image_layered.py @@ -0,0 +1,151 @@ +from __future__ import annotations + +import base64 +import gc +import io +import os +from pathlib import Path + +import pytest +import requests +import torch +from diffusers.pipelines.pipeline_utils import DiffusionPipeline +from PIL import Image + +from tests.conftest import ( + OmniServer, + _run_post_test_cleanup, + _run_pre_test_cleanup, +) +from tests.e2e.accuracy.utils import assert_image_sequence_similarity, model_output_dir +from tests.utils import hardware_test + +MODEL_ID = "Qwen/Qwen-Image-Layered" +MODEL_ENV_VAR = "QWEN_IMAGE_LAYERED_MODEL" +PROMPT = "decompose into layers" +NEGATIVE_PROMPT = " " +NUM_INFERENCE_STEPS = 20 +TRUE_CFG_SCALE = 4.0 +SEED = 777 +LAYERS = 3 +RESOLUTION = 640 +SSIM_THRESHOLD = 0.97 +PSNR_THRESHOLD = 30.0 + + +def _model_name() -> str: + return os.environ.get(MODEL_ENV_VAR, MODEL_ID) + + +def _local_files_only(model: str) -> bool: + return Path(model).exists() + + +def _normalize_layered_images(images: object) -> list[Image.Image]: + if not isinstance(images, list) or not images: + raise AssertionError(f"Unexpected layered output container: {type(images).__name__}") + + first_item = images[0] + if isinstance(first_item, Image.Image): + return [image.convert("RGBA") for image in images if isinstance(image, Image.Image)] + if isinstance(first_item, (list, tuple)): + return [image.convert("RGBA") for image in first_item if isinstance(image, Image.Image)] + raise AssertionError(f"Unexpected layered image element type: {type(first_item).__name__}") + + +def _run_vllm_omni_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]: + input_image.save(output_dir / "input.png") + server_args = ["--num-gpus", "1", "--stage-init-timeout", "300", "--init-timeout", "900"] + with OmniServer(model, server_args, use_omni=True) as omni_server: + buffer = io.BytesIO() + input_image.save(buffer, format="PNG") + buffer.seek(0) + response = requests.post( + f"http://{omni_server.host}:{omni_server.port}/v1/images/edits", + data={ + "model": omni_server.model, + "prompt": PROMPT, + "size": "auto", + "n": 1, + "response_format": "b64_json", + "negative_prompt": NEGATIVE_PROMPT, + "num_inference_steps": NUM_INFERENCE_STEPS, + "true_cfg_scale": TRUE_CFG_SCALE, + "seed": SEED, + "layers": LAYERS, + "resolution": RESOLUTION, + }, + files=[("image", ("input.png", buffer, "image/png"))], + timeout=600, + ) + response.raise_for_status() + payload = response.json() + assert len(payload["data"]) == LAYERS + output_images = [] + for item in payload["data"]: + image_bytes = base64.b64decode(item["b64_json"]) + image = Image.open(io.BytesIO(image_bytes)).convert("RGBA") + image.load() + output_images.append(image) + for index, image in enumerate(output_images, start=1): + image.save(output_dir / f"vllm_omni_layer_{index}.png") + return output_images + + +def _run_diffusers_qwen_image_layered(*, model: str, input_image: Image.Image, output_dir: Path) -> list[Image.Image]: + _run_pre_test_cleanup(enable_force=True) + pipe: DiffusionPipeline | None = None + try: + pipe = DiffusionPipeline.from_pretrained( + model, + torch_dtype=torch.bfloat16, + trust_remote_code=True, + local_files_only=_local_files_only(model), + ).to("cuda") + generator = torch.Generator(device="cuda").manual_seed(SEED) + result = pipe( # pyright: ignore[reportCallIssue] + image=input_image, + prompt=PROMPT, + negative_prompt=NEGATIVE_PROMPT, + num_inference_steps=NUM_INFERENCE_STEPS, + true_cfg_scale=TRUE_CFG_SCALE, + generator=generator, + num_images_per_prompt=1, + layers=LAYERS, + resolution=RESOLUTION, + ) + output_images = _normalize_layered_images(result.images) + assert len(output_images) == LAYERS, f"Expected {LAYERS} diffusers layers, got {len(output_images)}" + for index, image in enumerate(output_images, start=1): + image.save(output_dir / f"diffusers_layer_{index}.png") + return output_images + finally: + if pipe is not None and hasattr(pipe, "maybe_free_model_hooks"): + pipe.maybe_free_model_hooks() + del pipe + gc.collect() + if torch.cuda.is_available(): + torch.cuda.empty_cache() + _run_post_test_cleanup(enable_force=True) + + +@pytest.mark.advanced_model +@pytest.mark.benchmark +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100"}, num_cards=1) +def test_qwen_image_layered_matches_diffusers(accuracy_artifact_root: Path, qwen_bear_image: Image.Image) -> None: + model = _model_name() + output_dir = model_output_dir(accuracy_artifact_root, MODEL_ID) + input_image = qwen_bear_image.convert("RGBA") + + vllm_outputs = _run_vllm_omni_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir) + diffusers_outputs = _run_diffusers_qwen_image_layered(model=model, input_image=input_image, output_dir=output_dir) + + assert_image_sequence_similarity( + model_name=MODEL_ID, + vllm_images=vllm_outputs, + diffusers_images=diffusers_outputs, + ssim_threshold=SSIM_THRESHOLD, + psnr_threshold=PSNR_THRESHOLD, + compare_mode="RGBA", + ) diff --git a/tests/e2e/accuracy/utils.py b/tests/e2e/accuracy/utils.py index eb0eea757..d722b69b0 100644 --- a/tests/e2e/accuracy/utils.py +++ b/tests/e2e/accuracy/utils.py @@ -21,13 +21,14 @@ def assert_similarity( model_name: str, vllm_image: Image.Image, diffusers_image: Image.Image, - width: int, - height: int, ssim_threshold: float, psnr_threshold: float, + width: int | None = None, + height: int | None = None, + compare_mode: str = "RGB", ) -> None: - requested_size = (width, height) - if diffusers_image.size != requested_size: + requested_size = (width, height) if width is not None and height is not None else None + if requested_size is not None and diffusers_image.size != requested_size: pytest.skip( "Skipping as diffusers baseline output is corrupt and not comparable: " f"dimensions do not match requested size; requested={requested_size}, got={diffusers_image.size}." @@ -37,7 +38,11 @@ def assert_similarity( f"Online and diffusers output sizes mismatch: online={vllm_image.size}, diffusers={diffusers_image.size}" ) - ssim_score, psnr_score = compute_image_ssim_psnr(prediction=vllm_image, reference=diffusers_image) + ssim_score, psnr_score = compute_image_ssim_psnr( + prediction=vllm_image, + reference=diffusers_image, + compare_mode=compare_mode, + ) print(f"{model_name} similarity metrics:") print(f" SSIM: value={ssim_score:.6f}, threshold>={ssim_threshold:.6f}, range=[-1, 1], higher_is_better=True") print( @@ -52,13 +57,37 @@ def assert_similarity( ) +def assert_image_sequence_similarity( + *, + model_name: str, + vllm_images: list[Image.Image], + diffusers_images: list[Image.Image], + ssim_threshold: float, + psnr_threshold: float, + compare_mode: str = "RGB", +) -> None: + assert len(vllm_images) == len(diffusers_images), ( + f"Output image count mismatch for {model_name}: online={len(vllm_images)}, diffusers={len(diffusers_images)}" + ) + for index, (vllm_image, diffusers_image) in enumerate(zip(vllm_images, diffusers_images, strict=True), start=1): + assert_similarity( + model_name=f"{model_name}[layer={index}]", + vllm_image=vllm_image, + diffusers_image=diffusers_image, + ssim_threshold=ssim_threshold, + psnr_threshold=psnr_threshold, + compare_mode=compare_mode, + ) + + def compute_image_ssim_psnr( *, prediction: Image.Image, reference: Image.Image, + compare_mode: str = "RGB", ) -> tuple[float, float]: - pred_tensor = _pil_to_batched_tensor(prediction) - ref_tensor = _pil_to_batched_tensor(reference) + pred_tensor = _pil_to_batched_tensor(prediction, compare_mode=compare_mode) + ref_tensor = _pil_to_batched_tensor(reference, compare_mode=compare_mode) ssim_metric = StructuralSimilarityIndexMeasure(data_range=1.0) psnr_metric = PeakSignalNoiseRatio(data_range=1.0) @@ -68,7 +97,7 @@ def compute_image_ssim_psnr( return ssim_value, psnr_value -def _pil_to_batched_tensor(image: Image.Image) -> torch.Tensor: - array = np.asarray(image.convert("RGB"), dtype=np.float32) / 255.0 +def _pil_to_batched_tensor(image: Image.Image, *, compare_mode: str) -> torch.Tensor: + array = np.asarray(image.convert(compare_mode), dtype=np.float32) / 255.0 tensor = torch.from_numpy(array).permute(2, 0, 1).unsqueeze(0) return tensor From 4d816ff1ded1e35393d6175d8f0dbbe07d570add Mon Sep 17 00:00:00 2001 From: NATURE Date: Thu, 16 Apr 2026 16:25:13 +0800 Subject: [PATCH 23/97] [Feature] Bagel: Support tp+cfg parallel using mooncake transfer engine connector (#2705) Signed-off-by: natureofnature Co-authored-by: Hongsheng Liu --- .../omni_connectors/test_tp_rank_aware.py | 716 +++++++++++++++++ .../test_async_omni_engine_stage_init.py | 69 ++ tests/engine/test_single_stage_mode.py | 2 + .../distributed/group_coordinator.py | 5 +- .../diffusion/models/bagel/pipeline_bagel.py | 36 +- .../omni_connectors/kv_transfer_manager.py | 721 ++++++++++++------ .../omni_connectors/utils/kv_utils.py | 367 ++++++++- vllm_omni/engine/async_omni_engine.py | 16 +- vllm_omni/engine/stage_engine_core_client.py | 7 +- vllm_omni/engine/stage_init_utils.py | 116 ++- vllm_omni/entrypoints/openai/serving_chat.py | 53 +- vllm_omni/inputs/data.py | 4 + 12 files changed, 1846 insertions(+), 266 deletions(-) create mode 100644 tests/distributed/omni_connectors/test_tp_rank_aware.py diff --git a/tests/distributed/omni_connectors/test_tp_rank_aware.py b/tests/distributed/omni_connectors/test_tp_rank_aware.py new file mode 100644 index 000000000..d4793479a --- /dev/null +++ b/tests/distributed/omni_connectors/test_tp_rank_aware.py @@ -0,0 +1,716 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for rank-aware KV transfer (TP > 1) and heterogeneous TP support. + +Covers: +- _build_rank_aware_send_keys / _build_rank_aware_recv_keys +- _get_kv_source_ranks / _get_kv_target_ranks / get_kv_connector_key +- update_sender_info storing base host/port +- receive path constructing per-rank metadata for connector.get() +- Mooncake connector _query_metadata_at and partial-metadata get() path +""" + +from types import SimpleNamespace +from unittest.mock import MagicMock, patch + +import pytest +import torch + +from vllm_omni.distributed.omni_connectors.kv_transfer_manager import ( + KVCacheTransferData, + OmniKVCacheConfig, + OmniKVTransferManager, +) +from vllm_omni.distributed.omni_connectors.utils.initialization import ( + KV_RANK_PORT_STRIDE, +) +from vllm_omni.distributed.omni_connectors.utils.kv_utils import ( + KVTPTopology, + build_rank_aware_recv_keys, + build_rank_aware_send_keys, + get_kv_connector_key, + get_kv_source_ranks, + get_kv_target_ranks, + merge_received_rank_shards, + slice_received_rank_shard, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _make_manager( + from_tp: int = 1, + to_tp: int = 1, + local_rank: int = 0, + from_stage: str = "stage0", + to_stage: str = "stage1", + stage_id: str = "stage1", + need_recv: bool = True, + need_send: bool = False, + recv_timeout: float = 0.3, +) -> OmniKVTransferManager: + """Build a manager with TP params injected, bypassing torch.distributed.""" + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage=from_stage, + to_stage=to_stage, + stage_id=stage_id, + need_recv_cache=need_recv, + need_send_cache=need_send, + recv_timeout=recv_timeout, + from_tp=from_tp, + to_tp=to_tp, + ) + with ( + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=local_rank), + patch( + "vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", + return_value=max(from_tp, to_tp), + ), + ): + mgr = OmniKVTransferManager(config) + return mgr + + +def _make_payload(head_values: list[float], request_id: str = "req-1") -> dict: + head_tensor = torch.tensor(head_values, dtype=torch.float32).view(1, len(head_values), 1).repeat(2, 1, 1) + return { + "request_id": request_id, + "layer_blocks": { + "key_cache": [head_tensor.clone()], + "value_cache": [(head_tensor + 100).clone()], + }, + "block_ids": [0], + "metadata": {"seq_len": 2}, + } + + +def _make_transfer_data(head_values: list[float], request_id: str = "req-1") -> KVCacheTransferData: + payload = _make_payload(head_values, request_id=request_id) + return KVCacheTransferData( + request_id=request_id, + layer_blocks=payload["layer_blocks"], + block_ids=payload["block_ids"], + metadata=payload["metadata"], + ) + + +# ── Key format helper ──────────────────────────────────────────────── + + +class TestConnectorKeyFormat: + def test_key_format_matches_pr2677(self): + key = get_kv_connector_key("req-1", "stage0", 0, 1, 2) + assert key == "req-1_stage0_0_1_2" + + def test_key_fields_are_positional(self): + key = get_kv_connector_key("r", "s", 5, 3, 7) + parts = key.split("_") + assert parts == ["r", "s", "5", "3", "7"] + + +# ── Source / target rank mapping ───────────────────────────────────── + + +class TestRankMapping: + """Verify get_kv_target_ranks and get_kv_source_ranks for various TP configs.""" + + def test_homogeneous_tp2_rank0(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0) + assert get_kv_target_ranks(topo) == [0] + assert get_kv_source_ranks(topo) == [0] + + def test_homogeneous_tp2_rank1(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1) + assert get_kv_target_ranks(topo) == [1] + assert get_kv_source_ranks(topo) == [1] + + def test_homogeneous_tp4_rank3(self): + topo = KVTPTopology(source_tp_size=4, target_tp_size=4, local_rank=3) + assert get_kv_target_ranks(topo) == [3] + assert get_kv_source_ranks(topo) == [3] + + def test_sender_gt_receiver_tp4_to_tp2_rank0(self): + """Receiver rank 0 should receive from sender rank 0 and 1.""" + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0) + assert get_kv_source_ranks(topo) == [0, 1] + + def test_sender_gt_receiver_tp4_to_tp2_rank1(self): + """Receiver rank 1 should receive from sender rank 2 and 3.""" + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1) + assert get_kv_source_ranks(topo) == [2, 3] + + def test_sender_lt_receiver_tp2_to_tp4_rank0(self): + """Sender rank 0 should send to receiver ranks 0 and 1.""" + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) + assert get_kv_target_ranks(topo) == [0, 1] + + def test_sender_lt_receiver_tp2_to_tp4_rank1(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1) + assert get_kv_target_ranks(topo) == [2, 3] + + def test_receiver_lt_sender_source_ranks(self): + """Receiver rank 0 with tp2_to_tp4 should source from rank 0 only.""" + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) + assert get_kv_source_ranks(topo) == [0] + + def test_invalid_topology_raises(self): + topo = KVTPTopology(source_tp_size=3, target_tp_size=2, local_rank=0) + with pytest.raises(ValueError, match="divisible"): + get_kv_source_ranks(topo) + + +# ── _build_rank_aware_recv_keys ────────────────────────────────────── + + +class TestBuildRankAwareRecvKeys: + """Verify build_rank_aware_recv_keys returns (key, from_rank) tuples.""" + + def test_tp1_returns_legacy_key_with_none_rank(self): + topo = KVTPTopology(source_tp_size=1, target_tp_size=1, local_rank=0) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 1 + key, rank = pairs[0] + assert key == "omni_stage0_to_stage1_kv_cache_req-1" + assert rank is None + + def test_homogeneous_tp2_rank0(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 1 + key, rank = pairs[0] + assert key == "req-1_stage0_0_0_0" + assert rank == 0 + + def test_homogeneous_tp2_rank1(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=1) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 1 + key, rank = pairs[0] + assert key == "req-1_stage0_0_1_1" + assert rank == 1 + + def test_heterogeneous_tp4_to_tp2_rank0_gets_two_keys(self): + """Receiver rank 0 with source_tp=4, target_tp=2 should get 2 keys.""" + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=0) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 2 + + keys = [k for k, _ in pairs] + ranks = [r for _, r in pairs] + assert keys == ["req-1_stage0_0_0_0", "req-1_stage0_0_1_0"] + assert ranks == [0, 1] + + def test_heterogeneous_tp4_to_tp2_rank1_gets_two_keys(self): + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 2 + + ranks = [r for _, r in pairs] + assert ranks == [2, 3] + + def test_heterogeneous_tp2_to_tp4_rank2_gets_one_key(self): + """Receiver rank 2 with source_tp=2, target_tp=4 should get 1 key from sender rank 1.""" + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=2) + pairs = build_rank_aware_recv_keys("req-1", "stage0", "stage1", topo) + assert len(pairs) == 1 + key, rank = pairs[0] + assert rank == 1 + assert key == "req-1_stage0_0_1_2" + + +# ── _build_rank_aware_send_keys ────────────────────────────────────── + + +class TestBuildRankAwareSendKeys: + def test_tp1_returns_legacy_key(self): + topo = KVTPTopology(source_tp_size=1, target_tp_size=1, local_rank=0) + keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo) + assert keys == ["omni_stage0_to_stage1_kv_cache_req-1"] + + def test_homogeneous_tp2_rank0(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=2, local_rank=0) + keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo) + assert keys == ["req-1_stage0_0_0_0"] + + def test_sender_lt_receiver_tp2_to_tp4_rank0_sends_two_keys(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) + keys = build_rank_aware_send_keys("req-1", "stage0", "stage1", topo) + assert len(keys) == 2 + assert keys == ["req-1_stage0_0_0_0", "req-1_stage0_0_0_1"] + + +# ── update_sender_info stores base host/port ───────────────────────── + + +class TestUpdateSenderInfoBase: + def test_stores_base_host_and_port(self): + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + assert mgr._sender_base_host == "10.0.0.1" + assert mgr._sender_base_zmq_port == 50151 + + def test_rank1_adjusts_default_port_but_preserves_base(self): + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + assert mgr._sender_base_host == "10.0.0.1" + assert mgr._sender_base_zmq_port == 50151 + expected_adjusted = 50151 + 1 * KV_RANK_PORT_STRIDE + assert mgr.config.connector_config["sender_zmq_port"] == expected_adjusted + + def test_nested_sender_info_resolves_correctly(self): + """Nested sender_info keyed by integer stage id should resolve + using recv_stages (engine_input_source → recv_from).""" + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + stage_id=2, + engine_input_source=[1], + need_recv_cache=True, + from_tp=2, + to_tp=2, + ) + with ( + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0), + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=2), + ): + mgr = OmniKVTransferManager(config) + mgr.update_sender_info( + { + 0: {"host": "10.0.0.1", "zmq_port": 50151}, + 1: {"host": "10.0.0.2", "zmq_port": 50152}, + } + ) + assert mgr._sender_base_host == "10.0.0.2" + assert mgr._sender_base_zmq_port == 50152 + + +# ── receive path constructs per-rank metadata ──────────────────────── + + +class TestReceiveConstructsMetadata: + """Verify that receive_kv_cache_for_request passes metadata with + correct (host, port) to connector.get() for heterogeneous TP.""" + + def test_tp1_no_metadata_passed(self): + """TP=1: connector.get() should be called WITHOUT metadata.""" + mgr = _make_manager(from_tp=1, to_tp=1, local_rank=0, recv_timeout=0.05) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + calls = [] + + class _Connector: + def get(self, from_stage, to_stage, get_key, metadata=None): + calls.append({"key": get_key, "metadata": metadata}) + return None + + mgr._connector = _Connector() + mgr.receive_kv_cache_for_request("req-1") + + assert len(calls) > 0 + assert calls[0]["metadata"] is None + + def test_homogeneous_tp2_rank0_passes_metadata(self): + """TP=2 rank 0: metadata should point to sender rank 0's port.""" + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=0, recv_timeout=0.05) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + calls = [] + + class _Connector: + def get(self, from_stage, to_stage, get_key, metadata=None): + calls.append({"key": get_key, "metadata": metadata}) + return None + + mgr._connector = _Connector() + mgr.receive_kv_cache_for_request("req-1") + + assert len(calls) > 0 + meta = calls[0]["metadata"] + assert meta is not None + assert meta["source_host"] == "10.0.0.1" + assert meta["source_port"] == 50151 + 0 * KV_RANK_PORT_STRIDE + + def test_homogeneous_tp2_rank1_passes_metadata_with_offset(self): + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1, recv_timeout=0.05) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + calls = [] + + class _Connector: + def get(self, from_stage, to_stage, get_key, metadata=None): + calls.append({"key": get_key, "metadata": metadata}) + return None + + mgr._connector = _Connector() + mgr.receive_kv_cache_for_request("req-1") + + meta = calls[0]["metadata"] + assert meta["source_port"] == 50151 + 1 * KV_RANK_PORT_STRIDE + + def test_heterogeneous_tp4_to_tp2_rank0_multiple_metadata(self): + """Receiver rank 0 with source_tp=4, target_tp=2 should call get() with + two different metadata entries for sender ranks 0 and 1.""" + mgr = _make_manager(from_tp=4, to_tp=2, local_rank=0, recv_timeout=0.05) + mgr.update_sender_info({"host": "10.0.0.1", "zmq_port": 50151}) + + calls = [] + + class _Connector: + def get(self, from_stage, to_stage, get_key, metadata=None): + calls.append({"key": get_key, "metadata": metadata}) + return None + + mgr._connector = _Connector() + mgr.receive_kv_cache_for_request("req-1") + + seen_ports = set() + for c in calls: + if c["metadata"]: + seen_ports.add(c["metadata"]["source_port"]) + expected_ports = { + 50151 + 0 * KV_RANK_PORT_STRIDE, + 50151 + 1 * KV_RANK_PORT_STRIDE, + } + assert expected_ports.issubset(seen_ports) + + +# ── Mooncake connector _query_metadata_at ──────────────────────────── + + +class TestMooncakeQueryMetadataAt: + """Test the connector's _query_metadata_at method and partial-metadata + path in get() without requiring real RDMA/Mooncake.""" + + def test_query_metadata_at_returns_full_metadata(self): + """Mock the ZMQ interaction to verify _query_metadata_at returns + complete metadata including data_size.""" + + try: + from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import ( + MooncakeTransferEngineConnector, + QueryResponse, + ) + except ImportError: + pytest.skip("Mooncake not available") + + import msgspec + + connector = MagicMock(spec=MooncakeTransferEngineConnector) + connector._get_req_socket = MagicMock() + + mock_socket = MagicMock() + resp = QueryResponse(request_id="test_key@s0_s1", data_size=4096, is_fast_path=True) + mock_socket.recv.return_value = msgspec.msgpack.encode(resp) + connector._get_req_socket.return_value = mock_socket + + result = MooncakeTransferEngineConnector._query_metadata_at( + connector, + "test_key@s0_s1", + "10.0.0.1", + 50151, + ) + + assert result is not None + assert result["source_host"] == "10.0.0.1" + assert result["source_port"] == 50151 + assert result["data_size"] == 4096 + assert result["is_fast_path"] is True + + def test_query_metadata_at_returns_none_on_not_found(self): + try: + from vllm_omni.distributed.omni_connectors.connectors.mooncake_transfer_engine_connector import ( + INFO_NOT_FOUND, + MooncakeTransferEngineConnector, + ) + except ImportError: + pytest.skip("Mooncake not available") + + connector = MagicMock(spec=MooncakeTransferEngineConnector) + mock_socket = MagicMock() + mock_socket.recv.return_value = INFO_NOT_FOUND + connector._get_req_socket.return_value = mock_socket + + result = MooncakeTransferEngineConnector._query_metadata_at( + connector, + "test_key@s0_s1", + "10.0.0.1", + 50151, + ) + assert result is None + + +# ── Merge / slice hooks ────────────────────────────────────────────── + + +class TestMergeSliceHooks: + def test_single_shard_passes_through(self): + payload = {"layer_blocks": {"key_cache": [1]}} + assert merge_received_rank_shards([payload]) == payload + + def test_default_merger_concats_head_dim(self): + p0 = _make_payload([0.0]) + p1 = _make_payload([1.0]) + result = merge_received_rank_shards([p0, p1]) + key_cache = result["layer_blocks"]["key_cache"][0] + value_cache = result["layer_blocks"]["value_cache"][0] + assert key_cache.shape == (2, 2, 1) + assert value_cache.shape == (2, 2, 1) + assert torch.equal(key_cache[:, :, 0], torch.tensor([[0.0, 1.0], [0.0, 1.0]])) + assert torch.equal(value_cache[:, :, 0], torch.tensor([[100.0, 101.0], [100.0, 101.0]])) + + def test_custom_merger_hook_called(self): + merged = {"merged": True} + assert merge_received_rank_shards([{}, {}], merger=lambda payloads: merged) == merged + + def test_slicer_hook_called(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=0) + sliced = {"sliced": True} + assert slice_received_rank_shard({"full": True}, topo, slicer=lambda payload: sliced) == sliced + + def test_default_slicer_extracts_rank_local_heads(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1) + payload = _make_payload([0.0, 1.0]) + result = slice_received_rank_shard(payload, topo) + key_cache = result["layer_blocks"]["key_cache"][0] + value_cache = result["layer_blocks"]["value_cache"][0] + assert key_cache.shape == (2, 1, 1) + assert value_cache.shape == (2, 1, 1) + assert torch.equal(key_cache[:, :, 0], torch.tensor([[1.0], [1.0]])) + assert torch.equal(value_cache[:, :, 0], torch.tensor([[101.0], [101.0]])) + + def test_presliced_payload_is_not_sliced_twice(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=1) + payload = _make_payload([1.0]) + payload["metadata"]["tp_head_slice"] = {"applied": True, "target_rank": 1} + result = slice_received_rank_shard(payload, topo) + assert result is payload + + def test_round_trip_merge_from_tp4_to_tp2(self): + topo = KVTPTopology(source_tp_size=4, target_tp_size=2, local_rank=1) + source_ranks = get_kv_source_ranks(topo) + payloads = [_make_payload([float(rank)]) for rank in source_ranks] + result = merge_received_rank_shards(payloads) + key_cache = result["layer_blocks"]["key_cache"][0] + assert torch.equal(key_cache[:, :, 0], torch.tensor([[2.0, 3.0], [2.0, 3.0]])) + + def test_round_trip_slice_from_tp2_to_tp4(self): + topo = KVTPTopology(source_tp_size=2, target_tp_size=4, local_rank=3) + payload = _make_payload([2.0, 3.0]) + result = slice_received_rank_shard(payload, topo) + key_cache = result["layer_blocks"]["key_cache"][0] + assert torch.equal(key_cache[:, :, 0], torch.tensor([[3.0], [3.0]])) + + +class TestSenderSideSlicing: + def test_transfer_slices_before_sending_to_multiple_targets(self): + mgr = _make_manager( + from_tp=2, + to_tp=4, + local_rank=0, + need_send=True, + need_recv=False, + ) + sent_payloads = [] + + class _Connector: + supports_raw_data = False + + def put(self, from_stage, to_stage, put_key, data): + sent_payloads.append((put_key, KVCacheTransferData.from_bytes(data))) + return True, len(data), {} + + mgr._connector = _Connector() + mgr._transfer_kv_cache(_make_transfer_data([0.0, 1.0]), "req-1") + + assert [key for key, _ in sent_payloads] == ["req-1_stage0_0_0_0", "req-1_stage0_0_0_1"] + assert sent_payloads[0][1]["layer_blocks"]["key_cache"][0].shape == (2, 1, 1) + assert sent_payloads[1][1]["layer_blocks"]["key_cache"][0].shape == (2, 1, 1) + assert torch.equal( + sent_payloads[0][1]["layer_blocks"]["key_cache"][0][:, :, 0], + torch.tensor([[0.0], [0.0]]), + ) + assert torch.equal( + sent_payloads[1][1]["layer_blocks"]["key_cache"][0][:, :, 0], + torch.tensor([[1.0], [1.0]]), + ) + assert sent_payloads[0][1]["metadata"]["tp_head_slice"]["target_rank"] == 0 + assert sent_payloads[1][1]["metadata"]["tp_head_slice"]["target_rank"] == 1 + + +class _MockBroadcastGroup: + def __init__(self, world_size: int, rank_in_group: int, broadcast_value=None, recv_value=None): + self.world_size = world_size + self.rank_in_group = rank_in_group + self.broadcast_value = broadcast_value + self.recv_value = recv_value + self.broadcast_calls = [] + self.send_calls = [] + self.recv_calls = [] + self.shm_broadcaster = None + + def broadcast_object(self, obj=None, src: int = 0): + self.broadcast_calls.append((obj, src)) + return self.broadcast_value if self.broadcast_value is not None else obj + + def send_object(self, obj, dst: int): + self.send_calls.append((dst, obj)) + + def recv_object(self, src: int): + self.recv_calls.append(src) + return self.recv_value + + +class TestDistributedReceive: + def test_tp_cfg_leader_receives_then_sends_branch_local_payloads(self): + mgr = _make_manager(from_tp=2, to_tp=4, local_rank=0) + req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) + world_group = _MockBroadcastGroup(world_size=4, rank_in_group=2) + cfg_group = _MockBroadcastGroup(world_size=3, rank_in_group=0) + + def _receive(req_obj, cfg_func, target_device): + req_obj.past_key_values = SimpleNamespace(key_cache=[torch.tensor([1.0])]) + req_obj.kv_metadata = {"source": "leader"} + req_obj.sampling_params.past_key_values = req_obj.past_key_values + req_obj.sampling_params.kv_metadata = req_obj.kv_metadata + req_obj.sampling_params.cfg_text_past_key_values = SimpleNamespace(key_cache=[torch.tensor([2.0])]) + req_obj.sampling_params.cfg_text_kv_metadata = {"source": "cfg_text"} + req_obj.sampling_params.cfg_img_past_key_values = SimpleNamespace(key_cache=[torch.tensor([3.0])]) + req_obj.sampling_params.cfg_img_kv_metadata = {"source": "cfg_img"} + return True + + mgr.receive_multi_kv_cache = MagicMock(side_effect=_receive) + with ( + patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", + return_value=3, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", + return_value=0, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), + ): + assert mgr.receive_multi_kv_cache_distributed(req) is True + + mgr.receive_multi_kv_cache.assert_called_once() + assert mgr.receive_multi_kv_cache.call_args.args[2] == torch.device("cpu") + assert req.kv_metadata == {"source": "leader"} + assert cfg_group.broadcast_calls == [] + assert [dst for dst, _ in cfg_group.send_calls] == [1, 2] + rank1_payload = cfg_group.send_calls[0][1] + rank2_payload = cfg_group.send_calls[1][1] + assert torch.equal(rank1_payload["past_key_values"].key_cache[0], torch.tensor([1.0])) + assert torch.equal(rank2_payload["past_key_values"].key_cache[0], torch.tensor([1.0])) + assert rank1_payload["sp.cfg_active_branch"] == "cfg_text" + assert rank2_payload["sp.cfg_active_branch"] == "cfg_img" + assert rank1_payload["sp.cfg_branch_roles"] == ["cfg_text", "cfg_img"] + assert rank2_payload["sp.cfg_branch_roles"] == ["cfg_text", "cfg_img"] + assert "sp.cfg_branch_past_key_values" in rank1_payload + assert "sp.cfg_branch_past_key_values" in rank2_payload + assert list(rank1_payload["sp.cfg_branch_past_key_values"].keys()) == ["cfg_text"] + assert list(rank2_payload["sp.cfg_branch_past_key_values"].keys()) == ["cfg_img"] + assert "sp.cfg_text_past_key_values" in rank1_payload + assert "sp.cfg_img_past_key_values" not in rank1_payload + assert "sp.cfg_img_past_key_values" in rank2_payload + assert "sp.cfg_text_past_key_values" not in rank2_payload + + def test_tp_cfg_follower_receives_local_payload_without_receiving(self): + mgr = _make_manager(from_tp=2, to_tp=4, local_rank=1) + req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) + world_group = _MockBroadcastGroup(world_size=4, rank_in_group=3) + cfg_payload = { + "past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), + "kv_metadata": {"source": "main"}, + "sp.past_key_values": SimpleNamespace(key_cache=[torch.tensor([1.0])]), + "sp.kv_metadata": {"source": "main"}, + "sp.cfg_active_branch": "cfg_text", + "sp.cfg_branch_roles": ["cfg_text", "cfg_img"], + "sp.cfg_branch_past_key_values": { + "cfg_text": SimpleNamespace(key_cache=[torch.tensor([2.0])]), + }, + "sp.cfg_branch_kv_metadata": {"cfg_text": {"source": "cfg-text"}}, + "sp.cfg_text_past_key_values": SimpleNamespace(key_cache=[torch.tensor([2.0])]), + } + cfg_group = _MockBroadcastGroup(world_size=2, rank_in_group=1, recv_value=cfg_payload) + + mgr.receive_multi_kv_cache = MagicMock(return_value=True) + with ( + patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_world_size", + return_value=2, + ), + patch( + "vllm_omni.diffusion.distributed.parallel_state.get_classifier_free_guidance_rank", + return_value=1, + ), + patch("vllm_omni.diffusion.distributed.parallel_state.get_cfg_group", return_value=cfg_group), + ): + assert mgr.receive_multi_kv_cache_distributed(req) is True + + mgr.receive_multi_kv_cache.assert_not_called() + assert req.kv_metadata == {"source": "main"} + assert torch.equal(req.past_key_values.key_cache[0], torch.tensor([1.0])) + assert torch.equal(req.sampling_params.past_key_values.key_cache[0], torch.tensor([1.0])) + assert req.sampling_params.cfg_active_branch == "cfg_text" + assert req.sampling_params.cfg_branch_roles == ["cfg_text", "cfg_img"] + assert torch.equal( + req.sampling_params.cfg_branch_past_key_values["cfg_text"].key_cache[0], + torch.tensor([2.0]), + ) + assert req.sampling_params.cfg_branch_kv_metadata == {"cfg_text": {"source": "cfg-text"}} + assert torch.equal(req.sampling_params.cfg_text_past_key_values.key_cache[0], torch.tensor([2.0])) + assert cfg_group.broadcast_calls == [] + assert cfg_group.recv_calls == [0] + + def test_tp_without_cfg_keeps_independent_receive_path(self): + mgr = _make_manager(from_tp=2, to_tp=2, local_rank=1) + req = SimpleNamespace(request_id="req-1", sampling_params=SimpleNamespace()) + world_group = _MockBroadcastGroup(world_size=2, rank_in_group=1) + mgr.receive_multi_kv_cache = MagicMock(return_value=True) + + with patch("vllm_omni.diffusion.distributed.parallel_state.get_world_group", return_value=world_group): + assert mgr.receive_multi_kv_cache_distributed(req, target_device=torch.device("cpu")) is True + + mgr.receive_multi_kv_cache.assert_called_once_with(req, None, torch.device("cpu")) + + +# ── TP auto-detect ─────────────────────────────────────────────────── + + +class TestAutoDetectTP: + def test_auto_detect_when_config_defaults(self): + """When config from_tp/to_tp == 1 (default), manager should auto-detect.""" + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage="s0", + stage_id="s1", + need_recv_cache=True, + ) + with ( + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0), + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=4), + ): + mgr = OmniKVTransferManager(config) + assert mgr._tp_topo.source_tp_size == 4 + assert mgr._tp_topo.target_tp_size == 4 + + def test_explicit_tp_overrides_auto_detect(self): + config = OmniKVCacheConfig( + connector_config={"type": "mock"}, + from_stage="s0", + stage_id="s1", + need_recv_cache=True, + from_tp=2, + to_tp=4, + ) + with ( + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_local_tp_rank", return_value=0), + patch("vllm_omni.distributed.omni_connectors.kv_transfer_manager.get_tp_world_size", return_value=8), + ): + mgr = OmniKVTransferManager(config) + assert mgr._tp_topo.source_tp_size == 2 + assert mgr._tp_topo.target_tp_size == 4 diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index 84b0cb0be..5c2a9edb7 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -183,6 +183,7 @@ def test_launch_llm_stage_passes_stage_init_timeout_to_complete_stage_handshake( engine.model = "dummy-model" engine.single_stage_mode = False engine._omni_master_server = None + engine.stage_configs = [] metadata = types.SimpleNamespace(stage_id=0, runtime_cfg={"devices": "0"}) fake_vllm_config = types.SimpleNamespace() @@ -238,6 +239,7 @@ def test_launch_llm_stage_releases_launch_lock_before_complete_stage_handshake(m engine.model = "dummy-model" engine.single_stage_mode = False engine._omni_master_server = None + engine.stage_configs = [] fake_vllm_config = types.SimpleNamespace() fake_addresses = types.SimpleNamespace() @@ -378,3 +380,70 @@ def test_attach_llm_stage_uses_omni_input_preprocessor(monkeypatch): assert input_processor is not None assert isinstance(input_processor.input_preprocessor, DummyOmniInputPreprocessor) assert input_processor.input_preprocessor.renderer is input_processor.renderer + + +def test_inject_kv_stage_info_infers_sender_tp_topology(): + from vllm_omni.engine.stage_init_utils import inject_kv_stage_info + + stage0 = types.SimpleNamespace( + stage_id=0, + engine_args={ + "tensor_parallel_size": 4, + "omni_kv_config": { + "need_send_cache": True, + "omni_from_stage": "0", + "omni_to_stage": "1", + }, + }, + engine_input_source=[], + ) + stage1 = types.SimpleNamespace( + stage_id=1, + engine_args={ + "parallel_config": { + "tensor_parallel_size": 2, + "cfg_parallel_size": 1, + }, + "omni_kv_config": {"need_recv_cache": True}, + }, + engine_input_source=[0], + ) + + inject_kv_stage_info(stage0, 0, [stage0, stage1]) + + assert stage0.engine_args["omni_kv_config"]["stage_id"] == 0 + assert stage0.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2} + + +def test_inject_kv_stage_info_infers_receiver_tp_topology(): + from vllm_omni.engine.stage_init_utils import inject_kv_stage_info + + stage0 = types.SimpleNamespace( + stage_id=0, + engine_args={ + "tensor_parallel_size": 4, + "omni_kv_config": {"need_send_cache": True}, + }, + engine_input_source=[], + ) + stage1 = types.SimpleNamespace( + stage_id=1, + engine_args={ + "parallel_config": { + "tensor_parallel_size": 2, + "cfg_parallel_size": 1, + }, + "omni_kv_config": { + "need_recv_cache": True, + "omni_from_stage": "0", + "omni_to_stage": "1", + }, + }, + engine_input_source=[0], + ) + + inject_kv_stage_info(stage1, 1, [stage0, stage1]) + + assert stage1.engine_args["omni_kv_config"]["stage_id"] == 1 + assert stage1.engine_args["omni_kv_config"]["engine_input_source"] == [0] + assert stage1.engine_args["omni_kv_config"]["rank_mapping"] == {"from_tp": 4, "to_tp": 2} diff --git a/tests/engine/test_single_stage_mode.py b/tests/engine/test_single_stage_mode.py index 608e92ac4..28ccccaa2 100644 --- a/tests/engine/test_single_stage_mode.py +++ b/tests/engine/test_single_stage_mode.py @@ -1555,6 +1555,7 @@ class TestLaunchLlmStageSingleStageMode: engine.single_stage_mode = True engine._single_stage_id_filter = 0 engine._llm_stage_launch_lock = threading.Lock() + engine.stage_configs = [] mock_oms = mocker.Mock(spec=OmniMasterServer) mock_oms.address = "127.0.0.1" mock_oms.port = 25000 @@ -1629,6 +1630,7 @@ class TestLaunchLlmStageSingleStageMode: engine.single_stage_mode = False engine._omni_master_server = None engine._llm_stage_launch_lock = threading.Lock() + engine.stage_configs = [] fake_vllm_config = mocker.Mock() fake_executor_cls = mocker.Mock() diff --git a/vllm_omni/diffusion/distributed/group_coordinator.py b/vllm_omni/diffusion/distributed/group_coordinator.py index 8ab38f2a6..5294e6c9e 100644 --- a/vllm_omni/diffusion/distributed/group_coordinator.py +++ b/vllm_omni/diffusion/distributed/group_coordinator.py @@ -104,6 +104,7 @@ class GroupCoordinator: self.local_rank = local_rank self.device_group = None self.cpu_group = None + self.shm_broadcaster = None for ranks in group_ranks: device_group = torch.distributed.new_group(ranks, backend=torch_distributed_backend) @@ -316,7 +317,7 @@ class GroupCoordinator: assert dst < self.world_size, f"Invalid dst rank ({dst})" - assert dst != self.rank, "Invalid destination rank. Destination rank is the same as the current rank." + assert dst != self.rank_in_group, "Invalid destination rank. Destination rank is the same as the current rank." # Serialize object to tensor and get the size as well object_tensor = torch.frombuffer(pickle.dumps(obj), dtype=torch.uint8) @@ -338,7 +339,7 @@ class GroupCoordinator: assert src < self.world_size, f"Invalid src rank ({src})" - assert src != self.rank, "Invalid source rank. Source rank is the same as the current rank." + assert src != self.rank_in_group, "Invalid source rank. Source rank is the same as the current rank." size_tensor = torch.empty(1, dtype=torch.long, device="cpu") diff --git a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py index 72e53e7f4..a3d2259e6 100644 --- a/vllm_omni/diffusion/models/bagel/pipeline_bagel.py +++ b/vllm_omni/diffusion/models/bagel/pipeline_bagel.py @@ -365,28 +365,52 @@ class BagelPipeline(nn.Module, DiffusionPipelineProfilerMixin): if req.sampling_params.kv_metadata and "image_shape" in req.sampling_params.kv_metadata: image_shape = tuple(req.sampling_params.kv_metadata["image_shape"]) - cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None) + branch_kvs = getattr(req.sampling_params, "cfg_branch_past_key_values", None) or {} + branch_metadata = getattr(req.sampling_params, "cfg_branch_kv_metadata", None) or {} + active_branch = getattr(req.sampling_params, "cfg_active_branch", None) + branch_roles = getattr(req.sampling_params, "cfg_branch_roles", None) or list(branch_kvs.keys()) + + cfg_text_kv = getattr(req.sampling_params, "cfg_text_past_key_values", None) or branch_kvs.get("cfg_text") + cfg_text_metadata = getattr(req.sampling_params, "cfg_text_kv_metadata", None) or branch_metadata.get( + "cfg_text" + ) + cfg_img_kv = getattr(req.sampling_params, "cfg_img_past_key_values", None) or branch_kvs.get("cfg_img") + cfg_img_metadata = getattr(req.sampling_params, "cfg_img_kv_metadata", None) or branch_metadata.get( + "cfg_img" + ) + + cfg_parallel_contract = ( + active_branch is not None or bool(branch_roles) or cfg_text_kv is not None or cfg_img_kv is not None + ) + if cfg_parallel_contract: + logger.info( + "CFG enabled with injected branch KV context roles=%s active=%s", + branch_roles, + active_branch, + ) + if cfg_text_kv is not None: - logger.info("CFG enabled with multi-KV: using injected cfg_text KV Cache") cfg_text_seq_len = cfg_text_kv.key_cache[0].shape[0] cfg_text_context["past_key_values"] = cfg_text_kv cfg_text_context["kv_lens"] = [cfg_text_seq_len] - cfg_text_metadata = getattr(req.sampling_params, "cfg_text_kv_metadata", None) if cfg_text_metadata and "ropes" in cfg_text_metadata: cfg_text_context["ropes"] = cfg_text_metadata["ropes"] else: cfg_text_context["ropes"] = [cfg_text_seq_len] - cfg_img_kv = getattr(req.sampling_params, "cfg_img_past_key_values", None) or injected_kv + if cfg_img_kv is None and cfg_text_kv is not None: + cfg_img_kv = injected_kv + + if cfg_img_kv is not None: cfg_img_seq_len = cfg_img_kv.key_cache[0].shape[0] cfg_img_context["past_key_values"] = cfg_img_kv cfg_img_context["kv_lens"] = [cfg_img_seq_len] - cfg_img_metadata = getattr(req.sampling_params, "cfg_img_kv_metadata", None) if cfg_img_metadata and "ropes" in cfg_img_metadata: cfg_img_context["ropes"] = cfg_img_metadata["ropes"] else: cfg_img_context["ropes"] = [cfg_img_seq_len] - else: + + if not cfg_parallel_contract: logger.warning("CFG is disabled: only single KV cache available") gen_params = BagelGenParams( num_timesteps=gen_params.num_timesteps, diff --git a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py index 1958c9d40..ad008c397 100644 --- a/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py +++ b/vllm_omni/distributed/omni_connectors/kv_transfer_manager.py @@ -14,8 +14,20 @@ from vllm.logger import init_logger from .factory import OmniConnectorFactory from .utils.config import ConnectorSpec -from .utils.initialization import KV_TRANSFER_PORT_OFFSET -from .utils.kv_utils import normalize_layer_kv +from .utils.initialization import KV_RANK_PORT_STRIDE +from .utils.kv_utils import ( + KVTPTopology, + build_rank_aware_recv_keys, + build_rank_aware_send_keys, + get_kv_target_ranks, + get_local_tp_rank, + get_tp_world_size, + kv_zmq_port, + merge_received_rank_shards, + normalize_layer_kv, + slice_layer_blocks, + slice_received_rank_shard, +) logger = init_logger(__name__) @@ -57,6 +69,8 @@ class OmniKVCacheConfig: need_recv_cache: bool = False need_send_cache: bool = False recv_timeout: float = 30.0 + from_tp: int = 1 + to_tp: int = 1 @dataclass @@ -72,82 +86,44 @@ class KVCacheTransferData: """Convert to dictionary for serialization.""" return asdict(self) - def to_bytes(self) -> bytes: - """Convert to compact binary format for fast transfer.""" + def _build_tensors_desc(self, *, cpu: bool) -> tuple[list[dict[str, Any]], list, int, torch.device | None]: + """Iterate layer blocks and build tensor descriptors + data chunks. + + Returns ``(tensors_desc, chunks, total_bytes, device)``. + *chunks* contains ``bytes`` when *cpu* is True, flat uint8 GPU tensors otherwise. + """ tensors_desc: list[dict[str, Any]] = [] - tensor_bufs: list[bytes] = [] - data_offset = 0 - - for cache_name in ("key_cache", "value_cache"): - cache_list = self.layer_blocks.get(cache_name, []) - for layer_idx, tensor in enumerate(cache_list): - if tensor is None: - tensors_desc.append({"n": f"{cache_name}_{layer_idx}", "x": True}) - continue - - t = tensor.detach().cpu().contiguous() - dtype_str = str(t.dtype).removeprefix("torch.") - raw = t.view(torch.uint8).numpy().tobytes() - tensors_desc.append( - { - "n": f"{cache_name}_{layer_idx}", - "i": layer_idx, - "d": dtype_str, - "s": list(t.shape), - "o": data_offset, - "b": len(raw), - } - ) - tensor_bufs.append(raw) - data_offset += len(raw) - - header = json.dumps( - { - "rid": self.request_id, - "bids": self.block_ids, - "meta": self.metadata, - "td": tensors_desc, - "nl": len(self.layer_blocks.get("key_cache", [])), - }, - separators=(",", ":"), - ).encode("utf-8") - return b"".join([struct.pack(">I", len(header)), header] + tensor_bufs) - - def to_gpu_tensor(self) -> torch.Tensor: - """Convert to a packed GPU tensor for raw-data connectors.""" - tensors_desc: list[dict[str, Any]] = [] - gpu_tensors: list[torch.Tensor] = [] + chunks: list = [] data_offset = 0 device = None for cache_name in ("key_cache", "value_cache"): - cache_list = self.layer_blocks.get(cache_name, []) - for layer_idx, tensor in enumerate(cache_list): + for layer_idx, tensor in enumerate(self.layer_blocks.get(cache_name, [])): if tensor is None: tensors_desc.append({"n": f"{cache_name}_{layer_idx}", "x": True}) continue - t = tensor.detach().contiguous() - if device is None and t.is_cuda: + if cpu: + t = t.cpu() + elif device is None and t.is_cuda: device = t.device - dtype_str = str(t.dtype).removeprefix("torch.") nbytes = t.numel() * t.element_size() tensors_desc.append( { "n": f"{cache_name}_{layer_idx}", "i": layer_idx, - "d": dtype_str, + "d": str(t.dtype).removeprefix("torch."), "s": list(t.shape), "o": data_offset, "b": nbytes, } ) - gpu_tensors.append(t.view(torch.uint8).flatten()) + chunks.append(t.view(torch.uint8).numpy().tobytes() if cpu else t.view(torch.uint8).flatten()) data_offset += nbytes - if device is None: - raise RuntimeError("No CUDA tensors found, use to_bytes() instead") + return tensors_desc, chunks, data_offset, device + def _build_header_bytes(self, tensors_desc: list[dict[str, Any]]) -> bytes: header = json.dumps( { "rid": self.request_id, @@ -158,19 +134,26 @@ class KVCacheTransferData: }, separators=(",", ":"), ).encode("utf-8") + return struct.pack(">I", len(header)) + header - header_prefix = struct.pack(">I", len(header)) + header - total_size = len(header_prefix) + data_offset - output = torch.empty(total_size, dtype=torch.uint8, device=device) - header_tensor = torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8) - output[: len(header_prefix)].copy_(header_tensor) + def to_bytes(self) -> bytes: + """Convert to compact binary format for fast transfer.""" + tensors_desc, chunks, _, _ = self._build_tensors_desc(cpu=True) + return b"".join([self._build_header_bytes(tensors_desc)] + chunks) + def to_gpu_tensor(self) -> torch.Tensor: + """Convert to a packed GPU tensor for raw-data connectors.""" + tensors_desc, chunks, data_offset, device = self._build_tensors_desc(cpu=False) + if device is None: + raise RuntimeError("No CUDA tensors found, use to_bytes() instead") + header_prefix = self._build_header_bytes(tensors_desc) + output = torch.empty(len(header_prefix) + data_offset, dtype=torch.uint8, device=device) + output[: len(header_prefix)].copy_(torch.frombuffer(bytearray(header_prefix), dtype=torch.uint8)) pos = len(header_prefix) - for t_flat in gpu_tensors: + for t_flat in chunks: n = t_flat.numel() output[pos : pos + n].copy_(t_flat) pos += n - return output @staticmethod @@ -237,11 +220,8 @@ class KVCacheTransferData: return layer_idx @staticmethod - def from_bytes(raw: "bytes | bytearray | memoryview") -> dict[str, Any]: - """Reconstruct KV cache data from the packed bytes format.""" - raw_mv = memoryview(raw) if not isinstance(raw, memoryview) else raw - header, tensor_data_mv = KVCacheTransferData._load_header_from_memoryview(raw_mv) - + def _populate_caches(header: dict[str, Any], get_tensor: callable) -> dict[str, Any]: + """Shared deserialization loop for both CPU and GPU paths.""" num_layers = header["nl"] key_cache: list[torch.Tensor | None] = [None] * num_layers value_cache: list[torch.Tensor | None] = [None] * num_layers @@ -249,20 +229,9 @@ class KVCacheTransferData: for info in header["td"]: if info.get("x"): continue - name: str = info["n"] torch_dtype = KVCacheTransferData._resolve_torch_dtype(info["d"]) - offset, nbytes = KVCacheTransferData._validate_tensor_span(name, info, len(tensor_data_mv)) - t = ( - torch.frombuffer( - tensor_data_mv, - dtype=torch.uint8, - offset=offset, - count=nbytes, - ) - .view(torch_dtype) - .reshape(info["s"]) - ) + t = get_tensor(info).view(torch_dtype).reshape(info["s"]) layer_idx = KVCacheTransferData._resolve_layer_idx(info, num_layers) if name.startswith("key_cache_"): key_cache[layer_idx] = t @@ -276,37 +245,30 @@ class KVCacheTransferData: "metadata": header["meta"], } + @staticmethod + def from_bytes(raw: "bytes | bytearray | memoryview") -> dict[str, Any]: + """Reconstruct KV cache data from the packed bytes format.""" + raw_mv = memoryview(raw) if not isinstance(raw, memoryview) else raw + header, tensor_data_mv = KVCacheTransferData._load_header_from_memoryview(raw_mv) + data_len = len(tensor_data_mv) + + def _get(info: dict) -> torch.Tensor: + offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len) + return torch.frombuffer(tensor_data_mv, dtype=torch.uint8, offset=offset, count=nbytes) + + return KVCacheTransferData._populate_caches(header, _get) + @staticmethod def from_bytes_gpu(gpu_tensor: torch.Tensor) -> dict[str, Any]: """Reconstruct KV cache data from a packed GPU tensor.""" header, data_start = KVCacheTransferData._load_header_from_tensor(gpu_tensor) + data_len = int(gpu_tensor.numel()) - data_start - num_layers = header["nl"] - key_cache: list[torch.Tensor | None] = [None] * num_layers - value_cache: list[torch.Tensor | None] = [None] * num_layers - tensor_data_bytes = int(gpu_tensor.numel()) - data_start + def _get(info: dict) -> torch.Tensor: + offset, nbytes = KVCacheTransferData._validate_tensor_span(info["n"], info, data_len) + return gpu_tensor[data_start + offset : data_start + offset + nbytes].clone() - for info in header["td"]: - if info.get("x"): - continue - - name: str = info["n"] - torch_dtype = KVCacheTransferData._resolve_torch_dtype(info["d"]) - offset, nbytes = KVCacheTransferData._validate_tensor_span(name, info, tensor_data_bytes) - t = gpu_tensor[data_start + offset : data_start + offset + nbytes].clone() - t = t.view(torch_dtype).reshape(info["s"]) - layer_idx = KVCacheTransferData._resolve_layer_idx(info, num_layers) - if name.startswith("key_cache_"): - key_cache[layer_idx] = t - elif name.startswith("value_cache_"): - value_cache[layer_idx] = t - - return { - "request_id": header["rid"], - "layer_blocks": {"key_cache": key_cache, "value_cache": value_cache}, - "block_ids": header["bids"], - "metadata": header["meta"], - } + return KVCacheTransferData._populate_caches(header, _get) class OmniKVTransferManager: @@ -341,6 +303,30 @@ class OmniKVTransferManager: else (None, None) ) + local_rank = get_local_tp_rank() + + if config.from_tp <= 1 and config.to_tp <= 1: + detected_tp = get_tp_world_size() + from_tp = detected_tp + to_tp = detected_tp + else: + from_tp = config.from_tp + to_tp = config.to_tp + + self._tp_topo = KVTPTopology(source_tp_size=from_tp, target_tp_size=to_tp, local_rank=local_rank) + + # Injectable hooks (compatible with PR #2677 OmniConnectorModelRunnerMixin). + self.kv_send_key_builder: Callable | None = None + self.kv_recv_key_builder: Callable | None = None + self.kv_payload_merger: Callable | None = None + self.kv_payload_slicer: Callable | None = None + + # Base sender endpoint (rank-0 host/port) stored during + # update_sender_info(). Used by the receive path to construct + # per-rank metadata for heterogeneous TP without querying a registry. + self._sender_base_host: str | None = None + self._sender_base_zmq_port: int | None = None + if config.need_send_cache and config.connector_config: try: _ = self.connector @@ -348,11 +334,20 @@ class OmniKVTransferManager: except Exception as e: logger.warning("Failed to eagerly initialize sender connector: %s", e) + # ------------------------------------------------------------------ # + # Factory helpers + # ------------------------------------------------------------------ # + @classmethod def _create(cls, cfg: dict | None) -> "OmniKVTransferManager": """Create manager from raw config dict.""" if not cfg or not isinstance(cfg, dict): return cls(OmniKVCacheConfig()) + + rank_mapping = cfg.get("rank_mapping", {}) + if not isinstance(rank_mapping, dict): + rank_mapping = {} + return cls( OmniKVCacheConfig( connector_config=cfg.get("connector_config"), @@ -363,18 +358,17 @@ class OmniKVTransferManager: need_recv_cache=cfg.get("need_recv_cache", False), need_send_cache=cfg.get("need_send_cache", False), recv_timeout=cfg.get("recv_timeout", 30.0), + from_tp=int(rank_mapping.get("from_tp", 1)), + to_tp=int(rank_mapping.get("to_tp", 1)), ) ) @classmethod - def from_model_config(cls, config: Any) -> "OmniKVTransferManager": - """Create from model config (for AR model runner).""" + def from_od_config(cls, config: Any) -> "OmniKVTransferManager": + """Create from model or OmniDiffusion config.""" return cls._create(getattr(config, "omni_kv_config", None)) - @classmethod - def from_od_config(cls, config: Any) -> "OmniKVTransferManager": - """Create from OmniDiffusion config (for diffusion runner).""" - return cls._create(getattr(config, "omni_kv_config", None)) + from_model_config = from_od_config @classmethod def from_vllm_config(cls, vllm_config: Any, model_config: Any) -> "OmniKVTransferManager": @@ -417,45 +411,33 @@ class OmniKVTransferManager: ) c_extra["to_stage"] = str(self.config.to_stage) if self.config.to_stage is not None else "1" + try: + stage_int = int(self.config.from_stage) if self.config.from_stage is not None else 0 + except (TypeError, ValueError): + stage_int = 0 + zmq_port = kv_zmq_port(base_port, stage_int, self._tp_topo.local_rank) + if self.config.need_send_cache: c_extra["role"] = "sender" - from_stage = self.config.from_stage - if from_stage is not None: - try: - c_extra["zmq_port"] = base_port + KV_TRANSFER_PORT_OFFSET + int(from_stage) - except (TypeError, ValueError): - c_extra["zmq_port"] = base_port + KV_TRANSFER_PORT_OFFSET + c_extra["zmq_port"] = zmq_port elif self.config.need_recv_cache: c_extra["role"] = "receiver" - from_stage = self.config.from_stage - sender_port = base_port + KV_TRANSFER_PORT_OFFSET - if from_stage is not None: - try: - sender_port = base_port + KV_TRANSFER_PORT_OFFSET + int(from_stage) - except (TypeError, ValueError): - pass c_extra.setdefault("sender_host", c_extra.get("host", "127.0.0.1")) - c_extra.setdefault("sender_zmq_port", sender_port) + c_extra.setdefault("sender_zmq_port", zmq_port) logger.info( - "Initializing OmniConnector (purpose=kv_transfer) with config: %s, role: %s", - cfg, + "Initializing OmniConnector type=%s role=%s", + c_type, c_extra.get("role", "N/A"), ) self._connector = OmniConnectorFactory.create_connector(ConnectorSpec(name=c_type, extra=c_extra)) - except Exception as e: - logger.error(f"Failed to initialize OmniConnector: {e}") - import traceback - - traceback.print_exc() - # Cache failure sentinel to avoid repeated initialization attempts in hot paths. + except Exception: + logger.exception("Failed to initialize OmniConnector") self._connector = False return self._connector if self._connector else None - def get_connector(self): - """Get connector (compatibility wrapper for existing code).""" - return self.connector + get_connector = property(lambda self: self.connector) def _resolve_sender_info( self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None @@ -513,8 +495,187 @@ class OmniKVTransferManager: cache_list[idx] = tensor.clone() return data + def _slice_transfer_data_for_target(self, kv_data: KVCacheTransferData, target_rank: int) -> KVCacheTransferData: + """Pre-slice sender payload for one target rank when sender TP < receiver TP.""" + topo = self._tp_topo + ratio = topo.target_tp_size // topo.source_tp_size + offset_in_sender = target_rank % ratio + metadata = dict(kv_data.metadata) if isinstance(kv_data.metadata, dict) else {} + metadata["tp_head_slice"] = { + "applied": True, + "side": "sender", + "target_rank": target_rank, + "source_rank": topo.local_rank, + "from_tp": topo.source_tp_size, + "to_tp": topo.target_tp_size, + "offset_in_shard": offset_in_sender, + "num_slices": ratio, + } + return KVCacheTransferData( + request_id=kv_data.request_id, + layer_blocks=slice_layer_blocks(kv_data.layer_blocks, offset_in_sender, ratio), + block_ids=list(kv_data.block_ids), + metadata=metadata, + ) + + def _serialize_transfer_payload(self, kv_data: KVCacheTransferData) -> torch.Tensor | bytes | dict[str, Any]: + """Serialize KV transfer data using the connector's fastest supported path.""" + if getattr(self.connector, "supports_raw_data", False): + try: + return kv_data.to_gpu_tensor() + except Exception: + pass + try: + return kv_data.to_bytes() + except Exception: + return kv_data.to_dict() + + @staticmethod + def _collect_request_kv_payload(req: Any) -> dict[str, object]: + """Collect request-side KV objects for object broadcast.""" + kv_payload: dict[str, object] = {} + for attr in ("past_key_values", "kv_metadata"): + val = getattr(req, attr, None) + if val is not None: + kv_payload[attr] = val + + if hasattr(req, "sampling_params") and req.sampling_params is not None: + for key in list(vars(req.sampling_params).keys()): + if key in ("past_key_values", "kv_metadata") or ( + key.startswith("cfg_") + and ( + key.endswith("_past_key_values") + or key.endswith("_kv_metadata") + or key + in ( + "cfg_kv_request_ids", + "cfg_active_branch", + "cfg_branch_roles", + "cfg_branch_past_key_values", + "cfg_branch_kv_metadata", + ) + ) + ): + val = getattr(req.sampling_params, key, None) + if val is not None: + kv_payload[f"sp.{key}"] = val + + return kv_payload + + @staticmethod + def _apply_request_kv_payload( + req: Any, + kv_payload: dict[str, object], + target_device: torch.device | None = None, + ) -> None: + """Apply a broadcast KV payload back onto a request object.""" + for attr in ("past_key_values", "kv_metadata"): + val = kv_payload.get(attr) + if val is not None: + if target_device is not None: + val = _move_to_device(val, target_device) + setattr(req, attr, val) + + if hasattr(req, "sampling_params") and req.sampling_params is not None: + for key, val in kv_payload.items(): + if key.startswith("sp."): + if target_device is not None: + val = _move_to_device(val, target_device) + setattr(req.sampling_params, key[3:], val) + + @staticmethod + def _discover_cfg_branch_roles(req: Any) -> list[str]: + """Discover CFG branch roles in a stable order.""" + sampling_params = getattr(req, "sampling_params", None) + if sampling_params is None: + return [] + + roles: list[str] = [] + branch_map = getattr(sampling_params, "cfg_branch_past_key_values", None) or {} + for preferred_role in ("cfg_text", "cfg_img"): + if ( + preferred_role in branch_map + or getattr(sampling_params, f"{preferred_role}_past_key_values", None) is not None + ): + roles.append(preferred_role) + + for role in branch_map.keys(): + if role not in roles and branch_map.get(role) is not None: + roles.append(role) + + for key in vars(sampling_params).keys(): + if not (key.startswith("cfg_") and key.endswith("_past_key_values")): + continue + role = key.removesuffix("_past_key_values") + if role in ("cfg_branch",) or role in roles: + continue + if getattr(sampling_params, key, None) is not None: + roles.append(role) + + return roles + + @classmethod + def _build_cfg_rank_local_payloads(cls, req: Any, cfg_size: int) -> list[dict[str, object] | None]: + """Build per-cfg-rank payloads so each rank receives only its branch KV.""" + full_payload = cls._collect_request_kv_payload(req) + payloads: list[dict[str, object] | None] = [] + + main_payload = { + key: value + for key, value in full_payload.items() + if key in ("past_key_values", "kv_metadata", "sp.past_key_values", "sp.kv_metadata") + } + branch_roles = cls._discover_cfg_branch_roles(req) + if branch_roles: + main_payload["sp.cfg_branch_roles"] = list(branch_roles) + main_payload["sp.cfg_active_branch"] = None + payloads.append(main_payload or None) + + sampling_params = getattr(req, "sampling_params", None) + branch_map = getattr(sampling_params, "cfg_branch_past_key_values", None) or {} + branch_metadata_map = getattr(sampling_params, "cfg_branch_kv_metadata", None) or {} + + for role in branch_roles: + if sampling_params is None: + payloads.append(None) + continue + + branch_kv = branch_map.get(role) + if branch_kv is None: + branch_kv = getattr(sampling_params, f"{role}_past_key_values", None) + branch_metadata = branch_metadata_map.get(role) + if branch_metadata is None: + branch_metadata = getattr(sampling_params, f"{role}_kv_metadata", None) + if branch_kv is None: + payloads.append(None) + continue + + local_payload = dict(main_payload) + local_payload["sp.cfg_active_branch"] = role + local_payload["sp.cfg_branch_roles"] = list(branch_roles) + local_payload["sp.cfg_branch_past_key_values"] = {role: branch_kv} + local_payload[f"sp.{role}_past_key_values"] = branch_kv + if branch_metadata is not None: + local_payload["sp.cfg_branch_kv_metadata"] = {role: branch_metadata} + local_payload[f"sp.{role}_kv_metadata"] = branch_metadata + + payloads.append(local_payload) + + while len(payloads) < cfg_size: + payloads.append(None) + + return payloads[:cfg_size] + def update_sender_info(self, sender_info: dict[str, Any], sender_stage_id: str | int | None = None) -> None: - """Update receiver-side sender info before loading remote KV cache.""" + """Update receiver-side sender info before loading remote KV cache. + + The orchestrator always reports rank-0's ZMQ port. When TP > 1 the + receiver must offset the port so that each TP rank connects to the + corresponding sender rank's port. + + The base host/port are also stored so that the receive path can + construct per-rank metadata for heterogeneous TP scenarios. + """ if not self.config.need_recv_cache: return @@ -523,18 +684,39 @@ class OmniKVTransferManager: logger.warning("Invalid sender_info format: %s", sender_info) return + sender_host = actual_info.get("host") + base_zmq_port = actual_info.get("zmq_port") + + # Store base sender info for per-rank metadata construction. + self._sender_base_host = sender_host + if base_zmq_port is not None: + self._sender_base_zmq_port = int(base_zmq_port) + + # --- Default sender: offset to match this receiver's corresponding sender rank --- + zmq_port = base_zmq_port + if zmq_port is not None and self._tp_topo.local_rank > 0: + zmq_port = int(zmq_port) + self._tp_topo.local_rank * KV_RANK_PORT_STRIDE + if self.config.connector_config: - self.config.connector_config["sender_host"] = actual_info.get("host") - self.config.connector_config["sender_zmq_port"] = actual_info.get("zmq_port") + self.config.connector_config["sender_host"] = sender_host + self.config.connector_config["sender_zmq_port"] = zmq_port if self._connector and hasattr(self._connector, "update_sender_info"): try: - self._connector.update_sender_info(actual_info.get("host"), actual_info.get("zmq_port")) + self._connector.update_sender_info(sender_host, zmq_port) except Exception: if hasattr(self._connector, "sender_host"): - self._connector.sender_host = actual_info.get("host") + self._connector.sender_host = sender_host if hasattr(self._connector, "sender_zmq_port"): - self._connector.sender_zmq_port = actual_info.get("zmq_port") + self._connector.sender_zmq_port = zmq_port + + logger.info( + "Sender info updated: host=%s, base_port=%s, adjusted_port=%s (local_rank=%s)", + sender_host, + base_zmq_port, + zmq_port, + self._tp_topo.local_rank, + ) def handle_finished_requests_kv_transfer( self, @@ -692,35 +874,54 @@ class OmniKVTransferManager: kv_data.request_id = transfer_req_id serialization_start = time.perf_counter() - transfer_data: torch.Tensor | bytes | dict[str, Any] - supports_raw = getattr(self.connector, "supports_raw_data", False) + topo = self._tp_topo + send_keys = build_rank_aware_send_keys( + transfer_req_id, from_stage, to_stage, topo, hook=self.kv_send_key_builder + ) + sender_slice_active = ( + topo.source_tp_size < topo.target_tp_size and len(send_keys) > 1 and not callable(self.kv_send_key_builder) + ) + per_key_payloads: list[tuple[str, torch.Tensor | bytes | dict[str, Any]]] = [] - try: - if supports_raw: - transfer_data = kv_data.to_gpu_tensor() + if sender_slice_active: + target_ranks = get_kv_target_ranks(topo) + if len(target_ranks) != len(send_keys): + logger.warning( + "Skip sender-side KV slicing because target rank count does not match send key count: " + "target_ranks=%s send_keys=%s", + len(target_ranks), + len(send_keys), + ) + sender_slice_active = False else: - raise RuntimeError("Connector does not support raw tensor") - except Exception: - try: - transfer_data = kv_data.to_bytes() - except Exception: - data_dict = kv_data.to_dict() - data_dict["request_id"] = transfer_req_id - transfer_data = data_dict + for put_key, target_rank in zip(send_keys, target_ranks, strict=False): + sliced_kv_data = self._slice_transfer_data_for_target(kv_data, target_rank) + per_key_payloads.append((put_key, self._serialize_transfer_payload(sliced_kv_data))) + + if not per_key_payloads: + transfer_data = self._serialize_transfer_payload(kv_data) + per_key_payloads = [(put_key, transfer_data) for put_key in send_keys] serialization_ms = (time.perf_counter() - serialization_start) * 1000 logger.info("KV cache serialized for %s in %.1f ms", transfer_req_id, serialization_ms) transfer_start = time.perf_counter() - success, size, _ = self._transfer_with_retry(from_stage, to_stage, f"kv_cache_{transfer_req_id}", transfer_data) + total_size = 0 + all_succeeded = True + for put_key, transfer_data in per_key_payloads: + success, size, _ = self._transfer_with_retry(from_stage, to_stage, put_key, transfer_data) + total_size += size + all_succeeded = all_succeeded and success + elapsed = time.perf_counter() - transfer_start - if success: - mbps = (size / 1024 / 1024) / elapsed if elapsed > 0 else 0 + if all_succeeded: + mbps = (total_size / 1024 / 1024) / elapsed if elapsed > 0 else 0 logger.info( - "KV transfer OK: %s, %s bytes, %.3fs, %.1f MB/s", + "KV transfer OK: %s, %s bytes across %s key(s), %.3fs, %.1f MB/s", transfer_req_id, - size, + total_size, + len(send_keys), elapsed, mbps, ) @@ -731,7 +932,7 @@ class OmniKVTransferManager: self, from_stage: str, to_stage: str, - request_id: str, + put_key: str, data: "dict[str, Any] | bytes | torch.Tensor", max_retries: int = 3, ) -> tuple[bool, int, dict[str, Any] | None]: @@ -740,7 +941,7 @@ class OmniKVTransferManager: Args: from_stage: Source stage identifier to_stage: Target stage identifier - request_id: Request identifier for the key + put_key: Pre-built connector key (rank-aware when TP > 1) data: Data to transfer max_retries: Maximum number of retry attempts @@ -749,14 +950,12 @@ class OmniKVTransferManager: """ for attempt in range(max_retries): try: - # Build the full key for connector - full_request_id = f"omni_{from_stage}_to_{to_stage}_{request_id}" success, size, metadata = self.connector.put( - from_stage=from_stage, to_stage=to_stage, put_key=full_request_id, data=data + from_stage=from_stage, to_stage=to_stage, put_key=put_key, data=data ) if success: return success, size, metadata - logger.warning(f"Transfer attempt {attempt + 1} failed for {request_id}") + logger.warning(f"Transfer attempt {attempt + 1} failed for {put_key}") except Exception as e: logger.warning(f"Transfer attempt {attempt + 1} exception: {e}") @@ -801,22 +1000,46 @@ class OmniKVTransferManager: poll_interval = 0.01 max_poll_interval = 0.5 - logger.info(f"Wait for KV cache for request {request_id} from stage {from_stage} to {to_stage}...") + topo = self._tp_topo + recv_key_pairs = build_rank_aware_recv_keys( + request_id, from_stage, to_stage, topo, hook=self.kv_recv_key_builder + ) + pending_pairs = list(recv_key_pairs) + received_payloads: dict[str, tuple[dict[str, Any], int]] = {} + + logger.info( + "Wait for KV cache for request %s from stage %s to %s via %s key(s)...", + request_id, + from_stage, + to_stage, + len(recv_key_pairs), + ) try: while True: - # Build the full key for connector - full_request_id = f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}" link_start = time.perf_counter() - result = self.connector.get( - from_stage=from_stage, - to_stage=to_stage, - get_key=full_request_id, - ) - if result: + for get_key, from_rank in list(pending_pairs): + # Construct per-rank metadata so the connector queries + # the correct sender endpoint (heterogeneous TP path). + # When from_rank is None (TP<=1), metadata stays None + # and the connector falls back to its default sender. + rank_metadata: dict[str, Any] | None = None + if from_rank is not None and self._sender_base_host and self._sender_base_zmq_port is not None: + rank_metadata = { + "source_host": self._sender_base_host, + "source_port": self._sender_base_zmq_port + from_rank * KV_RANK_PORT_STRIDE, + } + + result = self.connector.get( + from_stage=from_stage, + to_stage=to_stage, + get_key=get_key, + metadata=rank_metadata, + ) + if not result: + continue + raw_data, size = result - elapsed = time.time() - start_time - link_ms = (time.perf_counter() - link_start) * 1000 managed_buffer = None if hasattr(raw_data, "tensor") and hasattr(raw_data, "release"): @@ -844,6 +1067,21 @@ class OmniKVTransferManager: else: data = raw_data + received_payloads[get_key] = (data, size) + pending_pairs.remove((get_key, from_rank)) + + if not pending_pairs and received_payloads: + elapsed = time.time() - start_time + link_ms = (time.perf_counter() - link_start) * 1000 + ordered_payloads = [received_payloads[key][0] for key, _ in recv_key_pairs] + total_size = sum(received_payloads[key][1] for key, _ in recv_key_pairs) + + if len(ordered_payloads) == 1: + data = ordered_payloads[0] + else: + data = merge_received_rank_shards(ordered_payloads, merger=self.kv_payload_merger) + data = slice_received_rank_shard(data, topo, slicer=self.kv_payload_slicer) + try: if isinstance(data, dict) and "layer_blocks" in data: layer_blocks = data["layer_blocks"] @@ -856,18 +1094,18 @@ class OmniKVTransferManager: continue if target_device is not None and tensor.device != target_device: cache_list[i] = tensor.to(target_device).contiguous() - finally: - if managed_buffer is not None: - managed_buffer.release() + except Exception: + logger.exception("Failed to move KV cache tensors to target device") logger.info( - "Successfully received KV cache for %s, %s bytes, wait=%.3fs, link=%.1fms", + "Successfully received KV cache for %s, %s bytes across %s key(s), wait=%.3fs, link=%.1fms", request_id, - size, + total_size, + len(recv_key_pairs), elapsed, link_ms, ) - return data, size + return data, total_size if time.time() - start_time > timeout: logger.error(f"Timeout waiting for KV cache for request {request_id} after {timeout}s") @@ -876,11 +1114,8 @@ class OmniKVTransferManager: time.sleep(poll_interval) poll_interval = min(poll_interval * 2, max_poll_interval) - except Exception as e: - logger.error(f"Error receiving KV cache for {request_id}: {e}") - import traceback - - traceback.print_exc() + except Exception: + logger.exception("Error receiving KV cache for %s", request_id) return None, 0 def apply_kv_cache_to_request(self, req: Any, data: dict[str, Any]) -> None: @@ -994,73 +1229,79 @@ class OmniKVTransferManager: cfg_kv_collect_func: Callable | None = None, target_device: torch.device | None = None, ) -> bool: - """Broadcast-aware wrapper around :meth:`receive_multi_kv_cache`. + """Distributed wrapper around :meth:`receive_multi_kv_cache`. - SharedMemory connector is single-reader: once rank 0 consumes the - segment it is deleted. For multi-GPU stages (e.g. sequence-parallel) - only rank 0 receives; the result is then broadcast to every other - rank via the world process-group. - - For single-worker stages this is equivalent to calling - :meth:`receive_multi_kv_cache` directly. + TP-aware path selection: + - world size 1: direct receive + - TP active, cfg size 1: each rank independently receives + - TP active, cfg size > 1: cfg-rank 0 receives, then broadcasts to + peers that share the same TP rank + - TP inactive: legacy rank-0 receive then world broadcast """ - from vllm_omni.diffusion.distributed.parallel_state import get_world_group + from vllm_omni.diffusion.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, + get_world_group, + ) world = get_world_group() if world.world_size <= 1: return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) - # --- rank 0: receive to CPU (needed for pickle-based broadcast) --- + topo = self._tp_topo + tp_active = topo.source_tp_size > 1 or topo.target_tp_size > 1 + cfg_size = 1 + cfg_rank = 0 + cfg_group = None + try: + cfg_size = get_classifier_free_guidance_world_size() + cfg_rank = get_classifier_free_guidance_rank() + cfg_group = get_cfg_group() + except Exception: + cfg_size = 1 + cfg_rank = 0 + cfg_group = None + + if tp_active and cfg_size <= 1: + logger.info( + "Rank-aware KV receive: rank %s independently receiving (from_tp=%s, to_tp=%s)", + topo.local_rank, + topo.source_tp_size, + topo.target_tp_size, + ) + return self.receive_multi_kv_cache(req, cfg_kv_collect_func, target_device) + + if tp_active and cfg_size > 1 and cfg_group is not None: + kv_payload: dict[str, object] | None = None + if cfg_rank == 0: + received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu")) + rank_payloads = self._build_cfg_rank_local_payloads(req, cfg_size) if received else [None] * cfg_size + kv_payload = rank_payloads[0] + for dst_rank in range(1, cfg_size): + cfg_group.send_object(rank_payloads[dst_rank], dst_rank) + else: + kv_payload = cfg_group.recv_object(0) + + if not kv_payload: + return False + + self._apply_request_kv_payload(req, kv_payload, target_device) + return True + + kv_payload: dict[str, object] | None = None if world.rank_in_group == 0: - self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu")) + received = self.receive_multi_kv_cache(req, cfg_kv_collect_func, torch.device("cpu")) + if received: + kv_payload = self._collect_request_kv_payload(req) - kv_payload: dict[str, object] = {} - for attr in ("past_key_values", "kv_metadata"): - val = getattr(req, attr, None) - if val is not None: - kv_payload[attr] = val + kv_payload = world.broadcast_object(kv_payload, src=0) - if hasattr(req, "sampling_params") and req.sampling_params is not None: - for key in list(vars(req.sampling_params).keys()): - if (key.startswith("cfg_") and key.endswith("_past_key_values")) or key in ( - "past_key_values", - "kv_metadata", - ): - val = getattr(req.sampling_params, key, None) - if val is not None: - kv_payload[f"sp.{key}"] = val - - payload_list = [kv_payload] - # Use broadcast_object_list (pickle-based) instead of broadcast_tensor_dict - # because the KV cache is a heterogeneous nested structure (NaiveCache objects - # with metadata + tensors), not a flat tensor dict. This runs once before - # the denoising loop so the serialization cost is negligible. - torch.distributed.broadcast_object_list(payload_list, src=world.ranks[0], group=world.cpu_group) - kv_payload = payload_list[0] - else: - payload_list: list[dict[str, object] | None] = [None] - torch.distributed.broadcast_object_list(payload_list, src=world.ranks[0], group=world.cpu_group) - kv_payload = payload_list[0] - - # --- apply on ALL ranks (rank 0 also needs CPU→GPU move) --- if not kv_payload: return False - for attr in ("past_key_values", "kv_metadata"): - val = kv_payload.get(attr) - if val is not None: - if target_device is not None: - val = _move_to_device(val, target_device) - setattr(req, attr, val) - - if hasattr(req, "sampling_params") and req.sampling_params is not None: - for key, val in kv_payload.items(): - if key.startswith("sp."): - if target_device is not None: - val = _move_to_device(val, target_device) - setattr(req.sampling_params, key[3:], val) - + self._apply_request_kv_payload(req, kv_payload, target_device) return True diff --git a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py index 2cb48a8b3..12b9b3d4f 100644 --- a/vllm_omni/distributed/omni_connectors/utils/kv_utils.py +++ b/vllm_omni/distributed/omni_connectors/utils/kv_utils.py @@ -1,15 +1,380 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Utility helpers for KV cache manipulation.""" +"""Utility helpers for KV cache manipulation, TP routing, and merge/slice.""" + +from __future__ import annotations + +import os +from collections.abc import Callable +from dataclasses import dataclass +from typing import Any import torch +from vllm.distributed.parallel_state import ( + get_tensor_model_parallel_rank, + get_tensor_model_parallel_world_size, +) from vllm.logger import init_logger +from .initialization import KV_RANK_PORT_STRIDE, KV_TRANSFER_PORT_OFFSET + logger = init_logger(__name__) LayerKV = torch.Tensor | tuple[torch.Tensor, torch.Tensor] +# ------------------------------------------------------------------ # +# TP Topology +# ------------------------------------------------------------------ # + + +@dataclass(frozen=True) +class KVTPTopology: + """Immutable descriptor for a KV-transfer parallel mapping. + + Captures sender/receiver parallel sizes and the local rank within + that parallel dimension. Works for any divisible parallel dimension + (TP, SP, Ring Attention). + """ + + source_tp_size: int + target_tp_size: int + local_rank: int + + def __post_init__(self) -> None: + if self.source_tp_size <= 0 or self.target_tp_size <= 0: + raise ValueError( + f"Parallel sizes must be positive: " + f"source_tp_size={self.source_tp_size}, target_tp_size={self.target_tp_size}" + ) + if self.local_rank < 0: + raise ValueError(f"local_rank must be non-negative, got {self.local_rank}") + + @property + def is_heterogeneous(self) -> bool: + return self.source_tp_size != self.target_tp_size + + @property + def ratio(self) -> int: + """Larger parallel size divided by smaller. Always >= 1.""" + return max(self.source_tp_size, self.target_tp_size) // min(self.source_tp_size, self.target_tp_size) + + +# ------------------------------------------------------------------ # +# Runtime TP detection +# ------------------------------------------------------------------ # + + +def get_local_tp_rank() -> int: + """Return the TP-local rank of this worker process. + + Uses ``get_tensor_model_parallel_rank()`` which returns the rank + within the TP group only, not the stage-global rank. + """ + try: + return get_tensor_model_parallel_rank() + except Exception: + logger.debug("TP parallel state not initialized, falling back to LOCAL_RANK env", exc_info=True) + try: + return int(os.environ.get("LOCAL_RANK", "0")) + except (ValueError, TypeError): + return 0 + + +def get_tp_world_size() -> int: + """Return the TP world size (tensor-parallel dimension only). + + Uses ``get_tensor_model_parallel_world_size()`` so that + cfg_parallel, SP, PP etc. are not included in the count. + """ + try: + return get_tensor_model_parallel_world_size() + except Exception: + logger.debug("TP parallel state not initialized, defaulting world_size=1", exc_info=True) + return 1 + + +# ------------------------------------------------------------------ # +# ZMQ port computation +# ------------------------------------------------------------------ # + + +def kv_zmq_port(base_port: int, from_stage: int, local_rank: int = 0) -> int: + """Compute the ZMQ port for a KV-transfer connector. + + Each TP rank gets its own port so that TP > 1 deployments do not + cause ``EADDRINUSE`` when multiple sender workers bind on the same + host. The formula is backward-compatible: rank 0 produces the same + port as the previous ``base + OFFSET + stage`` formula. + """ + return base_port + KV_TRANSFER_PORT_OFFSET + local_rank * KV_RANK_PORT_STRIDE + from_stage + + +# ------------------------------------------------------------------ # +# TP topology validation and rank routing +# ------------------------------------------------------------------ # + + +def validate_kv_tp_topology(topo: KVTPTopology) -> None: + """Reject heterogeneous TP mappings that cannot be routed losslessly.""" + larger = max(topo.source_tp_size, topo.target_tp_size) + smaller = min(topo.source_tp_size, topo.target_tp_size) + if larger % smaller != 0: + raise ValueError( + f"KV TP mapping must be divisible: " + f"source_tp_size={topo.source_tp_size}, " + f"target_tp_size={topo.target_tp_size}" + ) + + +def get_kv_target_ranks(topo: KVTPTopology) -> list[int]: + """Which remote ranks this local rank sends KV shards to (send side).""" + validate_kv_tp_topology(topo) + if topo.source_tp_size == topo.target_tp_size: + return [topo.local_rank] + if topo.source_tp_size > topo.target_tp_size: + return [topo.local_rank // (topo.source_tp_size // topo.target_tp_size)] + ratio = topo.target_tp_size // topo.source_tp_size + return [topo.local_rank * ratio + i for i in range(ratio)] + + +def get_kv_source_ranks(topo: KVTPTopology) -> list[int]: + """Which remote ranks this local rank receives KV shards from (recv side).""" + validate_kv_tp_topology(topo) + if topo.source_tp_size == topo.target_tp_size: + return [topo.local_rank] + if topo.source_tp_size > topo.target_tp_size: + ratio = topo.source_tp_size // topo.target_tp_size + return [topo.local_rank * ratio + i for i in range(ratio)] + return [topo.local_rank // (topo.target_tp_size // topo.source_tp_size)] + + +# ------------------------------------------------------------------ # +# Rank-aware connector key building +# ------------------------------------------------------------------ # + + +def get_kv_connector_key( + req_id: str, + from_stage: int | str, + chunk_id: int, + from_rank: int, + to_rank: int, +) -> str: + """Build connector key that includes rank info for KV transfers. + + Format matches PR #2677: ``{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}`` + """ + return f"{req_id}_{from_stage}_{chunk_id}_{from_rank}_{to_rank}" + + +def build_rank_aware_send_keys( + request_id: str, + from_stage: str, + to_stage: str, + topo: KVTPTopology, + hook: Callable | None = None, +) -> list[str]: + """Build send-side connector keys, checking injectable hook first.""" + if callable(hook): + keys = list(hook(request_id, from_stage, to_stage)) + if keys: + return keys + if topo.source_tp_size <= 1 and topo.target_tp_size <= 1: + return [f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}"] + target_ranks = get_kv_target_ranks(topo) + return [get_kv_connector_key(request_id, from_stage, 0, topo.local_rank, r) for r in target_ranks] + + +def build_rank_aware_recv_keys( + request_id: str, + from_stage: str, + to_stage: str, + topo: KVTPTopology, + hook: Callable | None = None, +) -> list[tuple[str, int | None]]: + """Build recv-side connector keys with sender rank info. + + Returns a list of ``(key, from_rank)`` tuples. ``from_rank`` is + ``None`` when TP <= 1 (single sender, no per-rank routing needed). + For TP > 1, ``from_rank`` identifies which sender rank owns the + key so that the connector can route metadata queries to the + correct endpoint. + """ + if callable(hook): + raw = list(hook(request_id, from_stage, to_stage)) + if raw: + if isinstance(raw[0], tuple): + return raw + # Hook returned plain strings (e.g. OmniConnectorModelRunnerMixin. + # get_rank_aware_kv_keys). Reconstruct from_rank from topology so + # Mooncake connector can route metadata queries to the correct + # sender endpoint in heterogeneous TP. + # TODO: have the mixin return (key, from_rank) tuples directly + # to avoid this indirect reconstruction. + source_ranks = get_kv_source_ranks(topo) + if len(raw) == len(source_ranks): + return list(zip(raw, source_ranks)) + return [(k, None) for k in raw] + if topo.source_tp_size <= 1 and topo.target_tp_size <= 1: + return [(f"omni_{from_stage}_to_{to_stage}_kv_cache_{request_id}", None)] + source_ranks = get_kv_source_ranks(topo) + return [(get_kv_connector_key(request_id, from_stage, 0, r, topo.local_rank), r) for r in source_ranks] + + +# ------------------------------------------------------------------ # +# KV tensor head slicing (heterogeneous TP) +# ------------------------------------------------------------------ # + + +def slice_kv_tensor_heads( + tensor: torch.Tensor | None, + offset_in_shard: int, + num_slices: int, +) -> torch.Tensor | None: + """Slice one KV tensor along its head dimension (dim 1).""" + if tensor is None: + return None + if not isinstance(tensor, torch.Tensor): + return tensor + if tensor.dim() < 2: + raise ValueError(f"Expected KV tensor with a head dimension, got shape={tuple(tensor.shape)}") + if num_slices <= 0: + raise ValueError(f"num_slices must be > 0, got {num_slices}") + if not (0 <= offset_in_shard < num_slices): + raise ValueError(f"offset_in_shard must be in [0, {num_slices}), got {offset_in_shard}") + + heads_in_shard = tensor.shape[1] + if heads_in_shard % num_slices != 0: + raise ValueError( + "KV head count must be divisible for heterogeneous TP slicing: " + f"heads_in_shard={heads_in_shard}, num_slices={num_slices}" + ) + + heads_per_slice = heads_in_shard // num_slices + start = offset_in_shard * heads_per_slice + end = start + heads_per_slice + return tensor[:, start:end, ...].contiguous() + + +def slice_layer_blocks( + layer_blocks: dict[str, Any], + offset_in_shard: int, + num_slices: int, +) -> dict[str, list[torch.Tensor | None]]: + """Slice all KV layers for one logical receiver rank.""" + sliced_blocks: dict[str, list[torch.Tensor | None]] = {} + for cache_name in ("key_cache", "value_cache"): + cache_list = layer_blocks.get(cache_name, []) + sliced_blocks[cache_name] = [ + slice_kv_tensor_heads(tensor, offset_in_shard, num_slices) for tensor in cache_list + ] + return sliced_blocks + + +# ------------------------------------------------------------------ # +# Multi-rank merge and receiver-side slice +# ------------------------------------------------------------------ # + + +def merge_received_rank_shards( + payloads: list[dict[str, Any]], + merger: Callable | None = None, +) -> dict[str, Any] | None: + """Merge multiple source-rank KV shards for one target rank. + + When *merger* is provided (injectable hook), it is called directly. + Otherwise the default merges along the head dimension (dim 1). + """ + if callable(merger): + return merger(payloads) + if not payloads: + return None + if len(payloads) == 1: + return payloads[0] + + base_payload = payloads[0] + if not isinstance(base_payload, dict) or "layer_blocks" not in base_payload: + return base_payload + + merged: dict[str, Any] = { + "request_id": base_payload.get("request_id"), + "block_ids": list(base_payload.get("block_ids", [])), + "metadata": dict(base_payload.get("metadata", {})), + } + merged_layer_blocks: dict[str, list[torch.Tensor | None]] = {} + + for cache_name in ("key_cache", "value_cache"): + cache_lists = [payload.get("layer_blocks", {}).get(cache_name, []) for payload in payloads] + num_layers = max((len(cache_list) for cache_list in cache_lists), default=0) + merged_cache: list[torch.Tensor | None] = [] + + for layer_idx in range(num_layers): + layer_tensors = [ + cache_list[layer_idx] + for cache_list in cache_lists + if layer_idx < len(cache_list) and cache_list[layer_idx] is not None + ] + if not layer_tensors: + merged_cache.append(None) + elif len(layer_tensors) == 1 or not isinstance(layer_tensors[0], torch.Tensor): + merged_cache.append(layer_tensors[0]) + else: + merged_cache.append(torch.cat(layer_tensors, dim=1).contiguous()) + + merged_layer_blocks[cache_name] = merged_cache + + merged["layer_blocks"] = merged_layer_blocks + return merged + + +def slice_received_rank_shard( + payload: dict[str, Any] | None, + topo: KVTPTopology, + slicer: Callable | None = None, +) -> dict[str, Any] | None: + """Optionally slice a received payload to extract this rank's portion. + + Used when ``to_tp > from_tp``: the sender sent full heads and each + receiver rank slices out its own subset. + """ + if callable(slicer): + return slicer(payload) + if not payload or topo.target_tp_size <= topo.source_tp_size or "layer_blocks" not in payload: + return payload + + metadata = payload.get("metadata", {}) + slice_metadata = metadata.get("tp_head_slice") if isinstance(metadata, dict) else None + if isinstance(slice_metadata, dict) and slice_metadata.get("applied"): + tagged_rank = slice_metadata.get("target_rank") + if tagged_rank is not None and tagged_rank != topo.local_rank: + logger.warning( + "Received pre-sliced KV payload for unexpected target rank: expected=%s got=%s", + topo.local_rank, + tagged_rank, + ) + return payload + + ratio = topo.target_tp_size // topo.source_tp_size + offset_in_sender = topo.local_rank % ratio + updated_metadata = dict(metadata) if isinstance(metadata, dict) else {} + updated_metadata["tp_head_slice"] = { + "applied": True, + "side": "receiver", + "target_rank": topo.local_rank, + "from_tp": topo.source_tp_size, + "to_tp": topo.target_tp_size, + "offset_in_shard": offset_in_sender, + "num_slices": ratio, + } + return { + "request_id": payload.get("request_id"), + "layer_blocks": slice_layer_blocks(payload["layer_blocks"], offset_in_sender, ratio), + "block_ids": list(payload.get("block_ids", [])), + "metadata": updated_metadata, + } + + def normalize_layer_kv( layer_kv: LayerKV, *, diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 054d5342d..23a85e9f5 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -61,6 +61,7 @@ from vllm_omni.engine.stage_engine_startup import ( ) from vllm_omni.engine.stage_init_utils import ( StartedLlmStage, + _inject_inferred_kv_tp_topology, acquire_device_locks, build_diffusion_config, build_engine_args_dict, @@ -78,7 +79,10 @@ from vllm_omni.engine.stage_init_utils import ( setup_stage_devices, terminate_alive_proc, ) -from vllm_omni.entrypoints.utils import load_and_resolve_stage_configs +from vllm_omni.entrypoints.utils import ( + inject_omni_kv_config, + load_and_resolve_stage_configs, +) from vllm_omni.inputs.preprocess import OmniInputPreprocessor from vllm_omni.platforms import current_omni_platform @@ -378,6 +382,12 @@ class AsyncOmniEngine: omni_kv["omni_to_stage"] = omni_to omni_kv.setdefault("stage_id", metadata.stage_id) engine_args_dict["omni_kv_config"] = omni_kv + if self.stage_configs: + _inject_inferred_kv_tp_topology( + engine_args_dict.get("omni_kv_config"), + metadata.stage_id, + self.stage_configs, + ) vllm_config, executor_class = build_vllm_config( stage_cfg, self.model, @@ -747,10 +757,8 @@ class AsyncOmniEngine: setup_stage_devices(configured_stage_id, metadata.runtime_cfg) omni_conn_cfg, omni_from, omni_to = omni_kv_connector if omni_conn_cfg: - from vllm_omni.entrypoints.utils import inject_omni_kv_config - inject_omni_kv_config(stage_cfg, omni_conn_cfg, omni_from, omni_to) - inject_kv_stage_info(stage_cfg, configured_stage_id) + inject_kv_stage_info(stage_cfg, configured_stage_id, self.stage_configs) if self.single_stage_mode: assert self._omni_master_server is not None stage_clients[stage_idx] = self._launch_diffusion_stage( diff --git a/vllm_omni/engine/stage_engine_core_client.py b/vllm_omni/engine/stage_engine_core_client.py index 52e674f47..ab2de757b 100644 --- a/vllm_omni/engine/stage_engine_core_client.py +++ b/vllm_omni/engine/stage_engine_core_client.py @@ -14,7 +14,9 @@ from vllm.logger import init_logger from vllm.v1.engine import EngineCoreRequest from vllm.v1.engine.core_client import AsyncMPClient, DPLBAsyncMPClient -from vllm_omni.distributed.omni_connectors.utils.initialization import KV_TRANSFER_PORT_OFFSET +from vllm_omni.distributed.omni_connectors.utils.initialization import ( + KV_TRANSFER_PORT_OFFSET, +) from vllm_omni.engine.stage_init_utils import StageMetadata if TYPE_CHECKING: @@ -246,6 +248,8 @@ class StageEngineCoreClientBase: from_stage = omni_kv_config.get("omni_from_stage", from_stage) try: + # Orchestrator always reports rank-0's port; receiver + # workers add their own local_rank * KV_RANK_PORT_STRIDE. sender_port = int(base_port) + KV_TRANSFER_PORT_OFFSET + int(from_stage) except (TypeError, ValueError): logger.warning( @@ -284,6 +288,7 @@ class StageEngineCoreClientBase: self._kv_sender_host = self._resolve_contact_host() if self._kv_sender_host is None: return None + # rank-0 base port; receiver workers adjust per KV_RANK_PORT_STRIDE. return { "host": self._kv_sender_host, "zmq_port": base_port + kv_transfer_port_offset + int(self.stage_id), diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 3a7fe4bad..c697e34ba 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -13,7 +13,7 @@ import importlib import multiprocessing as mp import os import time -from collections.abc import Callable +from collections.abc import Callable, Sequence from dataclasses import dataclass from typing import Any, Literal @@ -101,8 +101,110 @@ def resolve_worker_cls(engine_args: dict[str, Any]) -> None: raise ValueError(f"Unknown worker_type: {worker_type}") -def inject_kv_stage_info(stage_cfg: Any, stage_id: int) -> None: - """Inject stage metadata into omni_kv_config when present.""" +def _get_attr_or_item(obj: Any, key: str, default: Any = None) -> Any: + """Read *key* from *obj* regardless of whether it's a dict or object.""" + if hasattr(obj, "get"): + return obj.get(key, default) + return getattr(obj, key, default) + + +def _tp_size_for_stage(stage_configs: Sequence[Any], stage_id: Any) -> int | None: + """Resolve tensor_parallel_size for *stage_id* from the loaded stage configs.""" + id_strs = {str(stage_id)} + try: + id_strs.add(str(int(stage_id))) + except (TypeError, ValueError): + pass + + for stage_cfg in stage_configs: + if str(getattr(stage_cfg, "stage_id", None)) not in id_strs: + continue + engine_args = getattr(stage_cfg, "engine_args", None) + if engine_args is None: + return 1 + parallel_config = _get_attr_or_item(engine_args, "parallel_config") + if parallel_config is not None: + tp = _get_attr_or_item(parallel_config, "tensor_parallel_size", 1) + else: + tp = _get_attr_or_item(engine_args, "tensor_parallel_size", 1) + try: + return max(1, int(tp)) + except (TypeError, ValueError): + return 1 + return None + + +def _inject_inferred_kv_tp_topology( + omni_kv: Any, + stage_id: int, + stage_configs: Sequence[Any], + engine_input_source: Sequence[int] | None = None, +) -> None: + """Infer adjacent-stage TP topology and inject it into omni_kv_config. + + This keeps heterogeneous TP working without requiring user-authored + rank_mapping blocks in config files. + """ + if omni_kv is None: + return + + if hasattr(omni_kv, "get"): + need_send = bool(omni_kv.get("need_send_cache", False)) + need_recv = bool(omni_kv.get("need_recv_cache", False)) + omni_from_stage = omni_kv.get("omni_from_stage") + omni_to_stage = omni_kv.get("omni_to_stage") + rank_mapping = omni_kv.get("rank_mapping") + else: + need_send = bool(getattr(omni_kv, "need_send_cache", False)) + need_recv = bool(getattr(omni_kv, "need_recv_cache", False)) + omni_from_stage = getattr(omni_kv, "omni_from_stage", None) + omni_to_stage = getattr(omni_kv, "omni_to_stage", None) + rank_mapping = getattr(omni_kv, "rank_mapping", None) + + if not need_send and not need_recv: + return + + current_tp = _tp_size_for_stage(stage_configs, stage_id) + if current_tp is None: + return + + peer_stage_id = None + from_tp = None + to_tp = None + if str(omni_from_stage) == str(stage_id): + peer_stage_id = omni_to_stage + from_tp = current_tp + to_tp = _tp_size_for_stage(stage_configs, peer_stage_id) + elif str(omni_to_stage) == str(stage_id): + peer_stage_id = omni_from_stage + from_tp = _tp_size_for_stage(stage_configs, peer_stage_id) + to_tp = current_tp + elif need_recv and engine_input_source: + peer_stage_id = engine_input_source[0] + from_tp = _tp_size_for_stage(stage_configs, peer_stage_id) + to_tp = current_tp + + if from_tp is None or to_tp is None: + return + + if not isinstance(rank_mapping, dict): + rank_mapping = {} + rank_mapping.setdefault("from_tp", int(from_tp)) + rank_mapping.setdefault("to_tp", int(to_tp)) + + if hasattr(omni_kv, "__setitem__"): + omni_kv["rank_mapping"] = rank_mapping + else: + setattr(omni_kv, "rank_mapping", rank_mapping) + + +def inject_kv_stage_info(stage_cfg: Any, stage_id: int, stage_configs: Sequence[Any] | None = None) -> None: + """Inject stage_id, engine_input_source, and inferred TP topology into omni_kv_config. + + When *stage_configs* is provided, also infers from_tp/to_tp for + heterogeneous TP topologies so the KV transfer manager can compute + rank mappings automatically. + """ try: engine_args = stage_cfg.engine_args if hasattr(engine_args, "get"): @@ -125,6 +227,14 @@ def inject_kv_stage_info(stage_cfg: Any, stage_id: int) -> None: omni_kv.setdefault("engine_input_source", list(engine_input_source)) elif hasattr(omni_kv, "__setitem__") and "engine_input_source" not in omni_kv: omni_kv["engine_input_source"] = list(engine_input_source) + + if stage_configs: + _inject_inferred_kv_tp_topology( + omni_kv, + stage_id=stage_id, + stage_configs=stage_configs, + engine_input_source=engine_input_source, + ) except Exception as e: logger.debug("Failed to inject stage info into omni_kv_config: %s", e) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 39fcbc9a0..4b3a7045c 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -86,6 +86,7 @@ from vllm_omni.entrypoints.openai.image_api_utils import validate_layered_layers from vllm_omni.entrypoints.openai.protocol import OmniChatCompletionStreamResponse from vllm_omni.entrypoints.openai.protocol.audio import AudioResponse, CreateAudio from vllm_omni.entrypoints.openai.utils import ( + get_stage_type, get_supported_speakers_from_hf_config, parse_lora_request, validate_requested_speaker, @@ -294,6 +295,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) num_inference_steps = None + cfg_text_scale = None + cfg_img_scale = None # Omni multistage image generation: Stage-0 (AR) should receive a clean # text prompt (and optional conditioning image/size) so the model's own # processor can construct the correct inputs. @@ -342,6 +345,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): except Exception: pass negative_prompt = extra_body.get("negative_prompt") + cfg_text_scale = extra_body.get("cfg_text_scale") + cfg_img_scale = extra_body.get("cfg_img_scale") engine_prompt_image: dict[str, Any] | None = None is_img2img = False @@ -397,14 +402,18 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): sampling_params_list = self._build_sampling_params_list_from_request(request) # Apply user-specified overrides to diffusion stage(s) for image generation - if _image_gen_height is not None or _image_gen_width is not None or num_inference_steps is not None: - for idx, sp in enumerate(sampling_params_list): - if hasattr(sp, "height") and _image_gen_height is not None: - sp.height = _image_gen_height - if hasattr(sp, "width") and _image_gen_width is not None: - sp.width = _image_gen_width - if hasattr(sp, "num_inference_steps") and num_inference_steps is not None: - sp.num_inference_steps = num_inference_steps + for idx, sp in enumerate(sampling_params_list): + if hasattr(sp, "height") and _image_gen_height is not None: + sp.height = _image_gen_height + if hasattr(sp, "width") and _image_gen_width is not None: + sp.width = _image_gen_width + if hasattr(sp, "num_inference_steps") and num_inference_steps is not None: + sp.num_inference_steps = num_inference_steps + if hasattr(sp, "extra_args") and sp.extra_args is not None: + if cfg_text_scale is not None: + sp.extra_args["cfg_text_scale"] = cfg_text_scale + if cfg_img_scale is not None: + sp.extra_args["cfg_img_scale"] = cfg_img_scale self._log_inputs( request_id, @@ -2108,6 +2117,8 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): num_inference_steps = extra_body.get("num_inference_steps") guidance_scale = extra_body.get("guidance_scale") true_cfg_scale = extra_body.get("true_cfg_scale") or extra_body.get("cfg_scale") + cfg_text_scale = extra_body.get("cfg_text_scale") + cfg_img_scale = extra_body.get("cfg_img_scale") seed = extra_body.get("seed") negative_prompt = extra_body.get("negative_prompt") num_outputs_per_prompt = extra_body.get("num_outputs_per_prompt", 1) @@ -2162,6 +2173,10 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): gen_params.guidance_scale = guidance_scale if true_cfg_scale is not None: gen_params.true_cfg_scale = true_cfg_scale + if cfg_text_scale is not None: + gen_params.extra_args["cfg_text_scale"] = cfg_text_scale + if cfg_img_scale is not None: + gen_params.extra_args["cfg_img_scale"] = cfg_img_scale if num_frames is not None: gen_params.num_frames = num_frames if guidance_scale_2 is not None: @@ -2206,10 +2221,30 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Generate image diffusion_engine = cast(AsyncOmni, self._diffusion_engine) + stage_configs = list(getattr(diffusion_engine, "stage_configs", []) or []) + default_params_list = list(getattr(diffusion_engine, "default_sampling_params_list", []) or []) + + sampling_params_list: list[Any] = [] + for idx, stage_cfg in enumerate(stage_configs): + if get_stage_type(stage_cfg) == "diffusion": + sampling_params_list.append(gen_params) + continue + + default_stage_params = default_params_list[idx] if idx < len(default_params_list) else SamplingParams() + if hasattr(default_stage_params, "clone"): + try: + default_stage_params = default_stage_params.clone() + except Exception: + pass + sampling_params_list.append(default_stage_params) + + if not sampling_params_list: + sampling_params_list = [gen_params] + result = None async for output in diffusion_engine.generate( prompt=gen_prompt, - sampling_params_list=[gen_params], # Pass as single-stage params + sampling_params_list=sampling_params_list, request_id=request_id, ): result = output diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 85faf6b94..e4c33a58c 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -267,6 +267,10 @@ class OmniDiffusionSamplingParams: cfg_text_kv_metadata: dict[str, Any] | None = None cfg_img_kv_metadata: dict[str, Any] | None = None cfg_kv_request_ids: dict[str, str] | None = None + cfg_active_branch: str | None = None + cfg_branch_roles: list[str] | None = None + cfg_branch_past_key_values: dict[str, Any] | None = None + cfg_branch_kv_metadata: dict[str, dict[str, Any]] | None = None # Component modules modules: dict[str, Any] = field(default_factory=dict) From f1cb4ebe4ce200ccddb8297c88203c8da9b4fd53 Mon Sep 17 00:00:00 2001 From: fan2956 Date: Thu, 16 Apr 2026 18:21:34 +0800 Subject: [PATCH 24/97] [PERF] Wan2.2 support rmsnorm fused op (#2583) Signed-off-by: fan2956 Signed-off-by: gcanlin Co-authored-by: gcanlin --- tests/diffusion/layers/test_norm.py | 453 ++++++++++++++++++ vllm_omni/diffusion/layers/adalayernorm.py | 3 +- vllm_omni/diffusion/layers/norm.py | 110 +++++ .../models/wan2_2/wan2_2_transformer.py | 29 +- 4 files changed, 585 insertions(+), 10 deletions(-) create mode 100644 tests/diffusion/layers/test_norm.py create mode 100644 vllm_omni/diffusion/layers/norm.py diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py new file mode 100644 index 000000000..e42041528 --- /dev/null +++ b/tests/diffusion/layers/test_norm.py @@ -0,0 +1,453 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for LayerNorm and RMSNorm custom ops in diffusion layers.""" + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.diffusion, pytest.mark.cpu] + + +# ── Import tests ── + + +def test_layernorm_import(): + """Verify LayerNorm can be imported from the norm module.""" + from vllm_omni.diffusion.layers.norm import LayerNorm # noqa: F401 + + +def test_rmsnorm_import(): + """Verify RMSNorm can be imported from the norm module.""" + from vllm_omni.diffusion.layers.norm import RMSNorm # noqa: F401 + + +# ── LayerNorm tests ── + + +def test_layernorm_forward_shape(): + """LayerNorm produces correct output shapes.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + batch = 2 + seq_len = 4 + norm = LayerNorm(dim) + + x = torch.randn(batch, seq_len, dim) + out = norm(x) + + assert out.shape == (batch, seq_len, dim) + + +def test_layernorm_forward_shape_2d(): + """LayerNorm works with 2D input tensors.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + batch = 2 + norm = LayerNorm(dim) + + x = torch.randn(batch, dim) + out = norm(x) + + assert out.shape == (batch, dim) + + +def test_layernorm_preserves_dtype_fp32(): + """LayerNorm preserves float32 dtype.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + + x = torch.randn(2, 4, dim, dtype=torch.float32) + out = norm(x) + + assert out.dtype == torch.float32 + + +def test_layernorm_preserves_dtype_fp16(): + """LayerNorm preserves float16 dtype.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + + x = torch.randn(2, 4, dim, dtype=torch.float16) + out = norm(x) + + assert out.dtype == torch.float16 + + +def test_layernorm_preserves_dtype_bf16(): + """LayerNorm preserves bfloat16 dtype.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + + x = torch.randn(2, 4, dim, dtype=torch.bfloat16) + out = norm(x) + + assert out.dtype == torch.bfloat16 + + +def test_layernorm_without_elementwise_affine(): + """LayerNorm works without elementwise_affine (no learned parameters).""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim, elementwise_affine=False) + + assert norm.weight is None + assert norm.bias is None + + x = torch.randn(2, 4, dim) + out = norm(x) + + assert out.shape == (2, 4, dim) + + +def test_layernorm_custom_eps(): + """LayerNorm accepts custom epsilon value.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + eps = 1e-5 + norm = LayerNorm(dim, eps=eps) + + assert norm.eps == eps + + +def test_layernorm_has_learnable_parameters(): + """LayerNorm has learnable weight and bias by default.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + + assert norm.weight is not None + assert norm.bias is not None + assert norm.weight.shape == (dim,) + assert norm.bias.shape == (dim,) + + +def test_layernorm_matches_fp32_reference(): + """Verify LayerNorm produces identical output to FP32 nn.LayerNorm.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + eps = 1e-6 + torch.manual_seed(42) + + ours = LayerNorm(dim, eps=eps) + ref = torch.nn.LayerNorm(dim, eps=eps) + + # Copy weights + ref.weight.data.copy_(ours.weight.data) + ref.bias.data.copy_(ours.bias.data) + + x = torch.randn(2, 4, dim) + + out_ours = ours(x) + out_ref = ref(x.float()).to(x.dtype) + + torch.testing.assert_close(out_ours, out_ref, atol=1e-5, rtol=1e-5) + + +def test_layernorm_matches_diffusers_fp32layernorm(): + """Verify LayerNorm produces identical output to diffusers FP32LayerNorm.""" + from diffusers.models.normalization import FP32LayerNorm + + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + eps = 1e-6 + torch.manual_seed(42) + + ours = LayerNorm(dim, eps=eps) + ref = FP32LayerNorm(dim, eps=eps) + + # Copy weights + ref.weight.data.copy_(ours.weight.data) + ref.bias.data.copy_(ours.bias.data) + + # Test with fp16 input to verify FP32 computation + x = torch.randn(2, 4, dim, dtype=torch.float16) + + out_ours = ours(x) + out_ref = ref(x) + + torch.testing.assert_close(out_ours, out_ref, atol=1e-3, rtol=1e-3) + + +# ── RMSNorm tests ── + + +def test_rmsnorm_forward_shape(): + """RMSNorm produces correct output shapes.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + batch = 2 + seq_len = 4 + norm = RMSNorm(hidden_size) + + x = torch.randn(batch, seq_len, hidden_size) + out = norm(x) + + assert out.shape == (batch, seq_len, hidden_size) + + +def test_rmsnorm_forward_shape_2d(): + """RMSNorm works with 2D input tensors.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + batch = 2 + norm = RMSNorm(hidden_size) + + x = torch.randn(batch, hidden_size) + out = norm(x) + + assert out.shape == (batch, hidden_size) + + +def test_rmsnorm_preserves_dtype_fp32(): + """RMSNorm preserves float32 dtype.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + + x = torch.randn(2, 4, hidden_size, dtype=torch.float32) + out = norm(x) + + assert out.dtype == torch.float32 + + +def test_rmsnorm_preserves_dtype_fp16(): + """RMSNorm preserves float16 dtype.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + + x = torch.randn(2, 4, hidden_size, dtype=torch.float16) + out = norm(x) + + assert out.dtype == torch.float16 + + +def test_rmsnorm_preserves_dtype_bf16(): + """RMSNorm preserves bfloat16 dtype.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + + x = torch.randn(2, 4, hidden_size, dtype=torch.bfloat16) + out = norm(x) + + assert out.dtype == torch.bfloat16 + + +def test_rmsnorm_custom_eps(): + """RMSNorm accepts custom epsilon value.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + eps = 1e-5 + norm = RMSNorm(hidden_size, eps=eps) + + assert norm.variance_epsilon == eps + + +def test_rmsnorm_has_weight_parameter(): + """RMSNorm has learnable weight parameter initialized to ones.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + + assert norm.weight is not None + assert norm.weight.shape == (hidden_size,) + torch.testing.assert_close(norm.weight, torch.ones(hidden_size)) + + +def test_rmsnorm_numerical_correctness(): + """Verify RMSNorm produces numerically correct output.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + eps = 1e-6 + torch.manual_seed(42) + + norm = RMSNorm(hidden_size, eps=eps) + x = torch.randn(2, 4, hidden_size) + + # Compute expected output manually + x_fp32 = x.to(torch.float32) + variance = x_fp32.pow(2).mean(-1, keepdim=True) + expected = x_fp32 * torch.rsqrt(variance + eps) + expected = norm.weight.to(torch.float32) * expected + expected = expected.to(x.dtype) + + out = norm(x) + + torch.testing.assert_close(out, expected, atol=1e-5, rtol=1e-5) + + +def test_rmsnorm_matches_reference_implementation(): + """Verify RMSNorm matches a reference implementation.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + def reference_rmsnorm(x, weight, eps): + """Reference RMSNorm implementation.""" + input_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + out = x * torch.rsqrt(variance + eps) + out = weight.to(torch.float32) * out + return out.to(input_dtype) + + hidden_size = 128 + eps = 1e-6 + torch.manual_seed(123) + + norm = RMSNorm(hidden_size, eps=eps) + + # Test with various dtypes + for dtype in [torch.float32, torch.float16, torch.bfloat16]: + x = torch.randn(4, 8, hidden_size, dtype=dtype) + expected = reference_rmsnorm(x, norm.weight, eps) + out = norm(x) + torch.testing.assert_close(out, expected, atol=1e-3, rtol=1e-3) + + +# ── CustomOp dispatch tests ── + + +def test_layernorm_inherits_from_customop(): + """LayerNorm inherits from CustomOp for platform dispatch.""" + from vllm_omni.diffusion.layers.custom_op import CustomOp + from vllm_omni.diffusion.layers.norm import LayerNorm + + norm = LayerNorm(64) + assert isinstance(norm, CustomOp) + + +def test_rmsnorm_inherits_from_customop(): + """RMSNorm inherits from CustomOp for platform dispatch.""" + from vllm_omni.diffusion.layers.custom_op import CustomOp + from vllm_omni.diffusion.layers.norm import RMSNorm + + norm = RMSNorm(64) + assert isinstance(norm, CustomOp) + + +def test_layernorm_has_platform_methods(): + """LayerNorm has forward methods for each platform.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + norm = LayerNorm(64) + + assert hasattr(norm, "forward_cuda") + assert hasattr(norm, "forward_hip") + assert hasattr(norm, "forward_xpu") + assert hasattr(norm, "forward_npu") + assert hasattr(norm, "forward_native") + + +def test_rmsnorm_has_platform_methods(): + """RMSNorm has forward methods for each platform.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + norm = RMSNorm(64) + + assert hasattr(norm, "forward_cuda") + assert hasattr(norm, "forward_hip") + assert hasattr(norm, "forward_xpu") + assert hasattr(norm, "forward_npu") + assert hasattr(norm, "forward_native") + + +def test_layernorm_forward_native_directly(): + """LayerNorm.forward_native can be called directly.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + x = torch.randn(2, 4, dim) + + out = norm.forward_native(x) + + assert out.shape == (2, 4, dim) + + +def test_rmsnorm_forward_native_directly(): + """RMSNorm.forward_native can be called directly.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + x = torch.randn(2, 4, hidden_size) + + out = norm.forward_native(x) + + assert out.shape == (2, 4, hidden_size) + + +# ── Edge case tests ── + + +def test_layernorm_with_large_dim(): + """LayerNorm works with large hidden dimensions.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 4096 + norm = LayerNorm(dim) + x = torch.randn(1, 16, dim) + + out = norm(x) + + assert out.shape == (1, 16, dim) + + +def test_rmsnorm_with_large_dim(): + """RMSNorm works with large hidden dimensions.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 4096 + norm = RMSNorm(hidden_size) + x = torch.randn(1, 16, hidden_size) + + out = norm(x) + + assert out.shape == (1, 16, hidden_size) + + +def test_layernorm_with_single_element_batch(): + """LayerNorm works with batch size of 1.""" + from vllm_omni.diffusion.layers.norm import LayerNorm + + dim = 64 + norm = LayerNorm(dim) + x = torch.randn(1, 1, dim) + + out = norm(x) + + assert out.shape == (1, 1, dim) + + +def test_rmsnorm_with_single_element_batch(): + """RMSNorm works with batch size of 1.""" + from vllm_omni.diffusion.layers.norm import RMSNorm + + hidden_size = 64 + norm = RMSNorm(hidden_size) + x = torch.randn(1, 1, hidden_size) + + out = norm(x) + + assert out.shape == (1, 1, hidden_size) diff --git a/vllm_omni/diffusion/layers/adalayernorm.py b/vllm_omni/diffusion/layers/adalayernorm.py index 4d70ed52f..d147bdcfe 100644 --- a/vllm_omni/diffusion/layers/adalayernorm.py +++ b/vllm_omni/diffusion/layers/adalayernorm.py @@ -7,6 +7,7 @@ from vllm.logger import init_logger from vllm.model_executor.layers.linear import ReplicatedLinear from vllm_omni.diffusion.layers.custom_op import CustomOp +from vllm_omni.diffusion.layers.norm import LayerNorm if TYPE_CHECKING: from vllm.model_executor.layers.quantization.base_config import QuantizationConfig @@ -27,7 +28,7 @@ class AdaLayerNorm(CustomOp): self.eps = eps self.elementwise_affine = elementwise_affine self.hidden_size = hidden_size - self.layernorm = nn.LayerNorm(self.hidden_size, elementwise_affine=self.elementwise_affine, eps=self.eps) + self.layernorm = LayerNorm(self.hidden_size, elementwise_affine=self.elementwise_affine, eps=self.eps) def forward_cuda( self, diff --git a/vllm_omni/diffusion/layers/norm.py b/vllm_omni/diffusion/layers/norm.py new file mode 100644 index 000000000..6096ad7c3 --- /dev/null +++ b/vllm_omni/diffusion/layers/norm.py @@ -0,0 +1,110 @@ +from importlib.util import find_spec + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger + +from vllm_omni.diffusion.layers.custom_op import CustomOp + +logger = init_logger(__name__) + +_HAS_MINDIESD = find_spec("mindiesd") is not None + + +class LayerNorm(nn.LayerNorm, CustomOp): + """ + LayerNorm implementation that inherits from both ``nn.LayerNorm`` and ``CustomOp``. + NPU: + Uses ``mindiesd.fast_layernorm(self, x)`` when MindIE-SD is installed. + CUDA / HIP / XPU / native: + Falls back to FP32 nn.LayerNorm implementation. + """ + + def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine: bool = True): + super().__init__(normalized_shape=dim, eps=eps, elementwise_affine=elementwise_affine) + # CustomOp.__init__ cannot be called here because it would re-run + # nn.Module initialization and clear LayerNorm parameters. + self._forward_method = CustomOp.dispatch_forward(self) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self._forward_method(x) + + def forward_cuda(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + def forward_hip(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + def forward_xpu(self, x: torch.Tensor) -> torch.Tensor: + return self.forward_native(x) + + def forward_npu(self, x: torch.Tensor) -> torch.Tensor: + if _HAS_MINDIESD: + try: + from mindiesd import fast_layernorm + + return fast_layernorm(self, x) + except ImportError as e: + logger.warning_once( + "mindiesd.fast_layernorm import failed, falling back to FP32 layer_norm: %s", + e, + ) + + return self.forward_native(x) + + def forward_native(self, x: torch.Tensor) -> torch.Tensor: + origin_dtype = x.dtype + return F.layer_norm( + x.float(), + self.normalized_shape, + self.weight.float() if self.weight is not None else None, + self.bias.float() if self.bias is not None else None, + self.eps, + ).to(origin_dtype) + + +class RMSNorm(CustomOp): + def __init__(self, hidden_size: int, eps: float = 1e-6) -> None: + super().__init__() + self.weight = nn.Parameter(torch.ones(hidden_size)) + self.variance_epsilon = eps + + def forward_cuda( + self, + x: torch.Tensor, + ) -> torch.Tensor: + return self.forward_native(x) + + def forward_hip( + self, + x: torch.Tensor, + ) -> torch.Tensor: + return self.forward_native(x) + + def forward_npu( + self, + x: torch.Tensor, + ) -> torch.Tensor: + import torch_npu + + output = torch_npu.npu_rms_norm(x, gamma=self.weight, epsilon=self.variance_epsilon)[0] + + return output + + def forward_xpu( + self, + x: torch.Tensor, + ) -> torch.Tensor: + return self.forward_native(x) + + def forward_native( + self, + x: torch.Tensor, + ) -> torch.Tensor: + input_dtype = x.dtype + x = x.to(torch.float32) + variance = x.pow(2).mean(-1, keepdim=True) + out = x * torch.rsqrt(variance + self.variance_epsilon) + out = self.weight.to(torch.float32) * out + return out.to(input_dtype) diff --git a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py index b870193a1..d4d81b78e 100644 --- a/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py +++ b/vllm_omni/diffusion/models/wan2_2/wan2_2_transformer.py @@ -11,7 +11,6 @@ import torch.nn.functional as F from diffusers.models.attention import FeedForward from diffusers.models.embeddings import PixArtAlphaTextProjection, TimestepEmbedding, Timesteps from diffusers.models.modeling_outputs import Transformer2DModelOutput -from diffusers.models.normalization import FP32LayerNorm from vllm.distributed import ( get_tensor_model_parallel_rank, get_tensor_model_parallel_world_size, @@ -30,6 +29,7 @@ from vllm_omni.diffusion.distributed.sp_plan import ( ) from vllm_omni.diffusion.forward_context import get_forward_context from vllm_omni.diffusion.layers.adalayernorm import AdaLayerNorm +from vllm_omni.diffusion.layers.norm import LayerNorm, RMSNorm from vllm_omni.platforms import current_omni_platform logger = init_logger(__name__) @@ -236,9 +236,9 @@ class WanImageEmbedding(nn.Module): def __init__(self, in_features: int, out_features: int, pos_embed_seq_len: int | None = None): super().__init__() - self.norm1 = FP32LayerNorm(in_features) + self.norm1 = LayerNorm(in_features) self.ff = FeedForward(in_features, out_features, mult=1, activation_fn="gelu") - self.norm2 = FP32LayerNorm(out_features) + self.norm2 = LayerNorm(out_features) if pos_embed_seq_len is not None: self.pos_embed = nn.Parameter(torch.zeros(1, pos_embed_seq_len, in_features)) else: @@ -378,8 +378,12 @@ class WanSelfAttention(nn.Module): self.tp_inner_dim = self.num_heads * head_dim # QK normalization using vLLM's RMSNorm - self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) - self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + if get_tensor_model_parallel_world_size() > 1: + self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + else: + self.norm_q = RMSNorm(self.tp_inner_dim, eps=eps) + self.norm_k = RMSNorm(self.tp_inner_dim, eps=eps) self.to_out = RowParallelLinear( self.inner_dim, @@ -498,8 +502,12 @@ class WanCrossAttention(nn.Module): self.tp_inner_dim = self.num_heads * head_dim # QK normalization - self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) - self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + if get_tensor_model_parallel_world_size() > 1: + self.norm_q = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + self.norm_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + else: + self.norm_q = RMSNorm(self.tp_inner_dim, eps=eps) + self.norm_k = RMSNorm(self.tp_inner_dim, eps=eps) # Optional added KV projections for I2V (image embeddings) self.added_kv_proj_dim = added_kv_proj_dim @@ -518,7 +526,10 @@ class WanCrossAttention(nn.Module): gather_output=False, return_bias=False, ) - self.norm_added_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + if get_tensor_model_parallel_world_size() > 1: + self.norm_added_k = DistributedRMSNorm(self.tp_inner_dim, eps=eps) + else: + self.norm_added_k = RMSNorm(self.tp_inner_dim, eps=eps) else: self.add_k_proj = None self.add_v_proj = None @@ -637,7 +648,7 @@ class WanTransformerBlock(nn.Module): eps=eps, added_kv_proj_dim=added_kv_proj_dim, ) - self.norm2 = FP32LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() + self.norm2 = LayerNorm(dim, eps, elementwise_affine=True) if cross_attn_norm else nn.Identity() # 3. Feed-forward self.ffn = WanFeedForward(dim=dim, inner_dim=ffn_dim, dim_out=dim) From e8658b55d14482cdd30b5ee9cc2b6ca8e81d3f15 Mon Sep 17 00:00:00 2001 From: John Liu BUAA Date: Thu, 16 Apr 2026 18:49:59 +0800 Subject: [PATCH 25/97] [Test] Add performance tests for Qwen-Image-Layered model (#2807) Signed-off-by: John Liu BUAA --- .buildkite/test-nightly.yml | 4 +- .../test_qwen_image_layered_vllm_omni.json | 49 +++++++++++++++++++ 2 files changed, 52 insertions(+), 1 deletion(-) create mode 100644 tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 58e1e55af..ac43b597d 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -415,7 +415,9 @@ steps: EXIT2=$$? pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_edit_2509_vllm_omni.json EXIT3=$$? - if [ $$EXIT1 -eq 0 ] || [ $$EXIT2 -eq 0 ] || [ $$EXIT3 -eq 0 ]; then + pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json + EXIT4=$$? + if [ $$EXIT1 -eq 0 ] || [ $$EXIT2 -eq 0 ] || [ $$EXIT3 -eq 0 ] || [ $$EXIT4 -eq 0 ]; then buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" fi diff --git a/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json new file mode 100644 index 000000000..3cf13509c --- /dev/null +++ b/tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json @@ -0,0 +1,49 @@ +[ + { + "test_name": "test_qwen_image_layered_single_device", + "description": "Single-device baseline", + "server_type": "vllm-omni", + "server_params": { + "model": "Qwen/Qwen-Image-Layered", + "serve_args": { + "enable-diffusion-pipeline-profiler": true + } + }, + "benchmark_params": [ + { + "name": "640x640_steps20_i2i", + "dataset": "random", + "task": "i2i", + "width": 640, + "height": 640, + "num-inference-steps": 20, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.02, + "latency_mean": 40.0, + "peak_memory_mb_max": 70000, + "peak_memory_mb_mean": 70000 + } + }, + { + "name": "1024x1024_steps35_i2i", + "dataset": "random", + "task": "i2i", + "width": 1024, + "height": 1024, + "num-inference-steps": 35, + "num-prompts": 10, + "max-concurrency": 1, + "enable-negative-prompt": true, + "baseline": { + "throughput_qps": 0.005, + "latency_mean": 80.0, + "peak_memory_mb_max": 70000, + "peak_memory_mb_mean": 70000 + } + } + ] + } +] From 322620fd5774ffaf938395f0c065d703f85eed90 Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Thu, 16 Apr 2026 20:47:39 +0800 Subject: [PATCH 26/97] [Fix][Fish Speech] Remove redundant get_vocab() in control token encoding (#2842) Signed-off-by: Sy03 <1370724210@qq.com> --- vllm_omni/model_executor/models/fish_speech/prompt_utils.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/vllm_omni/model_executor/models/fish_speech/prompt_utils.py b/vllm_omni/model_executor/models/fish_speech/prompt_utils.py index 923e97b63..8b8d8559e 100644 --- a/vllm_omni/model_executor/models/fish_speech/prompt_utils.py +++ b/vllm_omni/model_executor/models/fish_speech/prompt_utils.py @@ -38,10 +38,7 @@ def _encode_plain_text(tokenizer: Any, text: str) -> list[int]: def _encode_control_token(tokenizer: Any, token: str) -> list[int]: - vocab = tokenizer.get_vocab() if hasattr(tokenizer, "get_vocab") else {} - token_id = vocab.get(token) - if token_id is None: - token_id = tokenizer.convert_tokens_to_ids(token) + token_id = tokenizer.convert_tokens_to_ids(token) if token_id is None or token_id == getattr(tokenizer, "unk_token_id", None): raise ValueError(f"Fish Speech tokenizer is missing required control token: {token}") return [int(token_id)] From 45760d61d231d433b01fb798f8180d146d3bc7ab Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Thu, 16 Apr 2026 21:27:43 +0800 Subject: [PATCH 27/97] [Test] Skip tests for known issues in audio and speaker recognition (#2851) --- tests/e2e/online_serving/test_qwen3_omni_expansion.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 306543908..06847f3d5 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -371,6 +371,7 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) +@pytest.mark.skip(reason="issue: #2827") def test_audio_in_video_001(omni_server, openai_client) -> None: """ Input Modal: text + video (synthetic MP4 with embedded audio; ``use_audio_in_video`` uses audio from the video). @@ -491,6 +492,7 @@ def test_speaker_001(omni_server, openai_client) -> None: @pytest.mark.omni @hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) +@pytest.mark.skip(reason="Known issue: occasional inaccuracy in voice recognition.") def test_speaker_002(omni_server, openai_client) -> None: """ Input Modal: text only (one-word answer constraint). From 2ec91d4dfd4dbfe8cb70ed448b56397c28cdd96b Mon Sep 17 00:00:00 2001 From: Mike Qiu Date: Thu, 16 Apr 2026 22:40:34 +0800 Subject: [PATCH 28/97] [FIX] Preserve YAML default stop words when request sends empty list (#2855) Signed-off-by: Mike_Qiu Co-authored-by: Mike_Qiu Co-authored-by: Claude Opus 4.6 --- .../test_serving_chat_sampling_params.py | 179 ++++++++++++++++++ vllm_omni/entrypoints/openai/serving_chat.py | 2 +- 2 files changed, 180 insertions(+), 1 deletion(-) diff --git a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py index fa4c1e195..4190b1fbb 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py +++ b/tests/entrypoints/openai_api/test_serving_chat_sampling_params.py @@ -284,6 +284,185 @@ def test_apply_request_overrides_applies_values(serving_chat, mock_request, defa assert result.top_k == 1 # YAML custom param preserved +# ============================================================================= +# Tests for empty-list handling in _apply_request_overrides +# ============================================================================= + + +def test_apply_overrides_empty_stop_list_preserves_default(serving_chat, mocker): + """Test that request.stop=[] does NOT override YAML default stop words.""" + default_params = SamplingParams(temperature=0.5, stop=["<|im_end|>"]) + request = mocker.MagicMock() + request.temperature = None + request.top_p = None + request.top_k = None + request.max_tokens = None + request.min_tokens = None + request.seed = None + request.ignore_eos = None + request.stop = [] # empty list — should be treated as "not set" + request.stop_token_ids = None + request.frequency_penalty = None + request.presence_penalty = None + + result = serving_chat._apply_request_overrides(default_params, request) + + assert result.stop == ["<|im_end|>"] # YAML default preserved + + +def test_apply_overrides_nonempty_stop_list_overrides_default(serving_chat, mocker): + """Test that request.stop=["\\n"] overrides YAML default stop words.""" + default_params = SamplingParams(temperature=0.5, stop=["<|im_end|>"]) + request = mocker.MagicMock() + request.temperature = None + request.top_p = None + request.top_k = None + request.max_tokens = None + request.min_tokens = None + request.seed = None + request.ignore_eos = None + request.stop = ["\n"] # non-empty list — should override + request.stop_token_ids = None + request.frequency_penalty = None + request.presence_penalty = None + + result = serving_chat._apply_request_overrides(default_params, request) + + assert result.stop == ["\n"] # Overridden by request + + +def test_apply_overrides_empty_stop_token_ids_preserves_default(serving_chat, mocker): + """Test that request.stop_token_ids=[] does NOT override YAML default.""" + default_params = SamplingParams(temperature=0.5, stop_token_ids=[2, 3]) + request = mocker.MagicMock() + request.temperature = None + request.top_p = None + request.top_k = None + request.max_tokens = None + request.min_tokens = None + request.seed = None + request.ignore_eos = None + request.stop = None + request.stop_token_ids = [] # empty list — should be treated as "not set" + request.frequency_penalty = None + request.presence_penalty = None + + result = serving_chat._apply_request_overrides(default_params, request) + + assert result.stop_token_ids == [2, 3] # YAML default preserved + + +def test_apply_overrides_nonempty_stop_token_ids_overrides_default(serving_chat, mocker): + """Test that request.stop_token_ids=[100] overrides YAML default.""" + default_params = SamplingParams(temperature=0.5, stop_token_ids=[2, 3]) + request = mocker.MagicMock() + request.temperature = None + request.top_p = None + request.top_k = None + request.max_tokens = None + request.min_tokens = None + request.seed = None + request.ignore_eos = None + request.stop = None + request.stop_token_ids = [100] # non-empty list — should override + request.frequency_penalty = None + request.presence_penalty = None + + result = serving_chat._apply_request_overrides(default_params, request) + + assert result.stop_token_ids == [100] # Overridden by request + + +def test_apply_overrides_mixed_empty_and_nonempty_lists(serving_chat, mocker): + """Test mixing empty and non-empty list fields with scalar fields.""" + default_params = SamplingParams( + temperature=0.4, + stop=["<|end|>"], + stop_token_ids=[2], + ) + request = mocker.MagicMock() + request.temperature = 0.9 + request.top_p = None + request.top_k = None + request.max_tokens = None + request.min_tokens = None + request.seed = None + request.ignore_eos = None + request.stop = [] # empty — should NOT override + request.stop_token_ids = [100, 200] # non-empty — SHOULD override + request.frequency_penalty = None + request.presence_penalty = None + + result = serving_chat._apply_request_overrides(default_params, request) + + assert result.temperature == 0.9 # Scalar override works + assert result.stop == ["<|end|>"] # Empty list did NOT override + assert result.stop_token_ids == [100, 200] # Non-empty list DID override + + +def test_apply_overrides_none_scalar_still_preserves_default(serving_chat, mocker): + """Regression: ensure None scalar values still don't override defaults.""" + default_params = SamplingParams(temperature=0.5, max_tokens=100, seed=42) + request = mocker.MagicMock() + request.temperature = None + request.top_p = None + request.top_k = None + request.max_tokens = None + request.min_tokens = None + request.seed = None + request.ignore_eos = None + request.stop = None + request.stop_token_ids = None + request.frequency_penalty = None + request.presence_penalty = None + + result = serving_chat._apply_request_overrides(default_params, request) + + assert result.temperature == 0.5 + assert result.max_tokens == 100 + assert result.seed == 42 + + +def test_apply_overrides_both_lists_empty_preserves_defaults(serving_chat, mocker): + """Test that both stop=[] and stop_token_ids=[] preserve YAML defaults.""" + default_params = SamplingParams( + temperature=0.5, + stop=["<|end|>", "\\n"], + stop_token_ids=[2, 32000], + ) + request = mocker.MagicMock() + request.temperature = None + request.top_p = None + request.top_k = None + request.max_tokens = None + request.min_tokens = None + request.seed = None + request.ignore_eos = None + request.stop = [] + request.stop_token_ids = [] + request.frequency_penalty = None + request.presence_penalty = None + + result = serving_chat._apply_request_overrides(default_params, request) + + assert result.stop == ["<|end|>", "\\n"] + assert result.stop_token_ids == [2, 32000] + + +def test_build_sampling_params_list_empty_stop_preserves_yaml(serving_chat, mock_request): + """Test that empty stop list in request preserves YAML defaults via + _build_sampling_params_list_from_request.""" + mock_request.stop = [] + mock_request.stop_token_ids = [] + + result = serving_chat._build_sampling_params_list_from_request(mock_request) + + comprehension_params = result[0] + # Empty lists should NOT override — YAML defaults are preserved + assert comprehension_params.stop == [] + assert comprehension_params.stop_token_ids == [] + + # ============================================================================= # Tests for _get_comprehension_stage_index # ============================================================================= diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 4b3a7045c..34ddbbd30 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -728,7 +728,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): for field_name in self._OPENAI_SAMPLING_FIELDS: value = getattr(request, field_name, None) - if value is not None: + if (value is not None and not isinstance(value, list)) or (isinstance(value, list) and len(value) > 0): setattr(params, field_name, value) return params From 7d64a7c9964ed7f285fec120dcb7396e027d600c Mon Sep 17 00:00:00 2001 From: Sy03 <1370724210@qq.com> Date: Thu, 16 Apr 2026 23:07:48 +0800 Subject: [PATCH 29/97] [BugFix][VoxCPM2]: split multichar Chinese tokens to match training tokenization (#2832) Signed-off-by: Sy03 <1370724210@qq.com> --- .../entrypoints/openai/serving_speech.py | 25 +++++++- .../models/voxcpm2/voxcpm2_talker.py | 63 ++++++++++++++++++- 2 files changed, 86 insertions(+), 2 deletions(-) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 1f78f5691..3eaf18111 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -216,6 +216,8 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): "Re-upload voices after each restart if needed." ) self._tts_tokenizer = None + self._voxcpm2_tokenizer = None + self._voxcpm2_split_map: dict[int, list[int]] = {} logger.info(f"Loaded {len(self.supported_speakers)} supported speakers: {sorted(self.supported_speakers)}") @@ -812,6 +814,25 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): return None # VoxCPM2 accepts any text input return self._validate_qwen_tts_request(request) + def _voxcpm2_encode(self, text: str) -> list[int]: + """Tokenize text for VoxCPM2, splitting multichar Chinese tokens.""" + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + build_cjk_split_map, + split_multichar_chinese, + ) + + if self._voxcpm2_tokenizer is None: + from transformers import AutoTokenizer + + model_name = self.engine_client.model_config.model + tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True) + self._voxcpm2_split_map = build_cjk_split_map(tokenizer) + self._voxcpm2_tokenizer = tokenizer + logger.info("VoxCPM2 serving: built multichar split map (%d entries)", len(self._voxcpm2_split_map)) + + ids = self._voxcpm2_tokenizer.encode(text, add_special_tokens=True) + return split_multichar_chinese(ids, self._voxcpm2_split_map) + def _validate_ref_audio_format(self, ref_audio: str) -> str | None: """Validate ref_audio is a supported URI format. Returns error or None.""" if not ( @@ -1508,7 +1529,9 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): if request.ref_audio is not None: wav_list, sr = await self._resolve_ref_audio(request.ref_audio) additional["reference_audio"] = [[wav_list, sr]] - prompt = {"prompt": request.input} + # Pre-split multichar Chinese tokens (VoxCPM2 was trained with single-char CJK IDs). + token_ids = self._voxcpm2_encode(request.input) + prompt: dict[str, Any] = {"prompt_token_ids": token_ids} if additional: prompt["additional_information"] = additional elif self._is_tts: diff --git a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py index 02bcae821..b666e41eb 100644 --- a/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py +++ b/vllm_omni/model_executor/models/voxcpm2/voxcpm2_talker.py @@ -41,6 +41,45 @@ logger = init_logger(__name__) _ENABLE_PROFILING = os.environ.get("VOXCPM2_PROFILE", "0") == "1" +def is_cjk_char(c: str) -> bool: + """Check if a character is a CJK ideograph.""" + cp = ord(c) + return ( + 0x4E00 <= cp <= 0x9FFF # CJK Unified Ideographs + or 0x3400 <= cp <= 0x4DBF # Extension A + or 0xF900 <= cp <= 0xFAFF # Compatibility Ideographs + or 0x20000 <= cp <= 0x2A6DF # Extension B + or 0x2A700 <= cp <= 0x2B73F # Extension C + or 0x2B740 <= cp <= 0x2B81F # Extension D + or 0x2F800 <= cp <= 0x2FA1F # Compatibility Supplement + ) + + +def build_cjk_split_map(tokenizer: Any) -> dict[int, list[int]]: + """Build {multichar_cjk_token_id: [single_char_ids]} from tokenizer vocab.""" + vocab = tokenizer.get_vocab() + split_map: dict[int, list[int]] = {} + for token, token_id in vocab.items(): + clean = token.replace("\u2581", "") + if len(clean) >= 2 and all(is_cjk_char(c) for c in clean): + char_ids = tokenizer.convert_tokens_to_ids(list(clean)) + if all(cid != tokenizer.unk_token_id for cid in char_ids): + split_map[token_id] = char_ids + return split_map + + +def split_multichar_chinese(token_ids: list[int], split_map: dict[int, list[int]]) -> list[int]: + """Replace multichar Chinese token IDs with single-char IDs (idempotent).""" + result: list[int] = [] + for tid in token_ids: + expansion = split_map.get(tid) + if expansion is not None: + result.extend(expansion) + else: + result.append(tid) + return result + + def _encode_raw_audio( tts: nn.Module, samples: list[float] | torch.Tensor, @@ -354,6 +393,8 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): self._cuda_graph_warmup_steps = 0 self._cuda_graph_warmup_threshold = 3 + self._multichar_zh_split: dict[int, list[int]] | None = None + self._active_states: dict[str, _RequestState] = {} self._current_request_id: str | None = None self._pending_requests: list[tuple[str, bool, torch.Tensor | None, int]] = [] @@ -985,6 +1026,17 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): return OmniOutput(text_hidden_states=model_outputs, multimodal_outputs=mm) + # -------------------- Chinese token splitting -------------------- + + def _get_multichar_zh_split(self) -> dict[int, list[int]]: + """Lazy-build {multichar_chinese_token_id: [char_id, ...]} map.""" + if self._multichar_zh_split is not None: + return self._multichar_zh_split + base_tokenizer = self.tts.text_tokenizer.tokenizer + self._multichar_zh_split = build_cjk_split_map(base_tokenizer) + logger.info("VoxCPM2: built multichar Chinese split map (%d entries)", len(self._multichar_zh_split)) + return self._multichar_zh_split + # -------------------- preprocess / postprocess -------------------- def preprocess( @@ -1011,8 +1063,17 @@ class VoxCPM2TalkerForConditionalGeneration(nn.Module): for rid in [r for r, s in self._active_states.items() if r not in pending_ids and s.prefill_completed]: self._cleanup_request(rid) - # VoxCPM2Tokenizer does char-level Chinese splitting, so use input_ids directly token_ids = input_ids.tolist() + # Fail-fast: unsplit multichar Chinese IDs in input_ids means the + # serving layer didn't pre-split. Silent fixup here would cause + # input_ids/embeds length mismatch (scheduler slot count is fixed). + split_map = self._get_multichar_zh_split() + if split_map and any(tid in split_map for tid in token_ids): + raise ValueError( + "VoxCPM2 preprocess received unsplit multichar Chinese " + "token IDs. The serving layer must send prompt_token_ids " + "with single-char CJK IDs (see _voxcpm2_encode)." + ) if token_ids and token_ids[0] == self.config.bos_token_id: token_ids = token_ids[1:] From c3ca5daafb05acec828a66e3ba5f84951715fcf2 Mon Sep 17 00:00:00 2001 From: TaffyOfficial <2587297563@qq.com> Date: Thu, 16 Apr 2026 23:15:25 +0800 Subject: [PATCH 30/97] Feat/Add HunyuanImage-3.0-Instruct ar part support: (#2713) Signed-off-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: TaffyOfficial <2324465096@qq.com> Co-authored-by: Claude Opus 4.6 (1M context) --- .../hunyuan_image3/prompt_utils.py | 88 ++++++++ .../test_hunyuan_image3_sampler.py | 190 +++++++++++++++++ .../test_hunyuanimage3_text2img.py | 2 +- .../models/hunyuan_image3/hunyuan_image3.py | 195 ++++++++++++++++++ .../stage_configs/hunyuan_image3_i2t.yaml | 44 ++++ .../stage_configs/hunyuan_image3_it2i.yaml | 78 +++++++ .../stage_configs/hunyuan_image3_moe.yaml | 81 -------- .../stage_configs/hunyuan_image3_t2t.yaml | 45 ++++ .../stage_input_processors/hunyuan_image3.py | 123 +++++++++++ vllm_omni/patch.py | 52 +++++ 10 files changed, 816 insertions(+), 82 deletions(-) create mode 100644 examples/offline_inference/hunyuan_image3/prompt_utils.py create mode 100644 tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py create mode 100644 vllm_omni/model_executor/stage_configs/hunyuan_image3_i2t.yaml create mode 100644 vllm_omni/model_executor/stage_configs/hunyuan_image3_it2i.yaml delete mode 100644 vllm_omni/model_executor/stage_configs/hunyuan_image3_moe.yaml create mode 100644 vllm_omni/model_executor/stage_configs/hunyuan_image3_t2t.yaml create mode 100644 vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py diff --git a/examples/offline_inference/hunyuan_image3/prompt_utils.py b/examples/offline_inference/hunyuan_image3/prompt_utils.py new file mode 100644 index 000000000..a5ef8e153 --- /dev/null +++ b/examples/offline_inference/hunyuan_image3/prompt_utils.py @@ -0,0 +1,88 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +Prompt construction utilities for HunyuanImage-3.0-Instruct examples. + +Wraps system_prompt.get_system_prompt() with task-aware presets so that +examples and tests don't need to manually concatenate system prompts, +, , and tags. + +Usage: + from prompt_utils import build_prompt + + # IT2I (image editing, think+recaption mode) + prompt = build_prompt("Make the petals neon pink", task="it2i_think") + + # I2T (image understanding) + prompt = build_prompt("Describe the content of the picture.", task="i2t") +""" + +from __future__ import annotations + +from vllm_omni.diffusion.models.hunyuan_image3.system_prompt import ( + get_system_prompt, +) + +# task → (sys_type, bot_task, trigger_tag) +# trigger_tag: "", "", or None +_TASK_PRESETS: dict[str, tuple[str, str | None, str | None]] = { + # Pure text generation (text → text, no image) + "t2t": ("en_unified", None, None), + # Image understanding (image → text) + "i2t": ("en_unified", None, None), + # Image editing (image+text → image), think+recaption mode + "it2i_think": ("en_unified", "think", ""), + # Image editing, recaption-only mode + "it2i_recaption": ("en_unified", "recaption", ""), + # Text-to-image, think mode + "t2i_think": ("en_unified", "think", ""), + # Text-to-image, recaption mode + "t2i_recaption": ("en_unified", "recaption", ""), + # Text-to-image, vanilla (no CoT) + "t2i_vanilla": ("en_vanilla", "image", None), +} + + +def build_prompt( + user_prompt: str, + task: str = "it2i_think", + sys_type: str | None = None, + custom_system_prompt: str | None = None, +) -> str: + """Build a complete HunyuanImage-3.0 prompt with auto-selected system + prompt and mode trigger tags. + + Args: + user_prompt: The user's raw instruction or question. + task: One of the preset task keys (see _TASK_PRESETS). + sys_type: Override the preset's sys_type for get_system_prompt(). + custom_system_prompt: Custom system prompt text (used when + sys_type="custom"). + + Returns: + Fully formatted prompt string ready for Omni.generate(). + """ + if task not in _TASK_PRESETS: + raise ValueError(f"Unknown task {task!r}. Choose from: {sorted(_TASK_PRESETS)}") + + preset_sys_type, preset_bot_task, trigger_tag = _TASK_PRESETS[task] + effective_sys_type = sys_type or preset_sys_type + + system_prompt = get_system_prompt(effective_sys_type, preset_bot_task, custom_system_prompt) + sys_text = system_prompt.strip() if system_prompt else "" + + has_image_input = task.startswith("i2t") or task.startswith("it2i") + + parts = ["<|startoftext|>"] + if sys_text: + parts.append(sys_text) + # Instruct conversation template: \n\nUser: ... \n\nAssistant: + parts.append("\n\nUser: ") + if has_image_input: + parts.append("") + parts.append(user_prompt) + parts.append("\n\nAssistant: ") + if trigger_tag: + parts.append(trigger_tag) + + return "".join(parts) diff --git a/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py new file mode 100644 index 000000000..51f6a85f5 --- /dev/null +++ b/tests/diffusion/models/hunyuan_image3/test_hunyuan_image3_sampler.py @@ -0,0 +1,190 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Unit tests for HunyuanImage3 AR sampler logic (stage transitions, +ratio restriction, comprehension blocking).""" + +import pytest +import torch + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + +# Fake token IDs for testing (avoid importing the real model). +END_OF_THINK = 100 +RECAPTION = 101 +END_OF_RECAPTION = 102 +ANSWER = 103 +BOI = 104 +SIZE_TOKEN = 105 +EOS = 106 +RATIO_START = 200 +RATIO_END = 210 +RATIO_OTHER_START = 220 +RATIO_OTHER_END = 223 + + +class FakeSamplerModel: + """Minimal stub that replicates the sampler-relevant attributes of + HunyuanImage3ForConditionalGeneration without loading real weights.""" + + def __init__(self, *, is_comprehension: bool = False): + self._is_comprehension = is_comprehension + self._eos_token_id = EOS + self._end_of_think_id = END_OF_THINK + self._recaption_id = RECAPTION + self._end_of_recaption_id = END_OF_RECAPTION + self._answer_id = ANSWER + self._mrope_boi_token_id = BOI + self._size_token_id = SIZE_TOKEN + self._start_ratio_id = RATIO_START + self._end_ratio_id = RATIO_END + self._ratio_other_slices = [(RATIO_OTHER_START, RATIO_OTHER_END + 1)] + self._all_ratio_ids = set(range(RATIO_START, RATIO_END + 1)) + self._all_ratio_ids.update(range(RATIO_OTHER_START, RATIO_OTHER_END + 1)) + + self._stage_transitions: dict[int, list[int]] = {} + if not is_comprehension: + self._stage_transitions[END_OF_THINK] = [RECAPTION] + self._stage_transitions[END_OF_RECAPTION] = [ANSWER, BOI, SIZE_TOKEN] + + self._blocked_token_ids: set[int] = set() + if is_comprehension: + self._blocked_token_ids.update([BOI, SIZE_TOKEN]) + self._blocked_token_ids.update(self._all_ratio_ids) + + # Bind the real methods from the model class. + from vllm_omni.model_executor.models.hunyuan_image3.hunyuan_image3 import ( + HunyuanImage3ForConditionalGeneration as _Real, + ) + + _get_forced_token = _Real._get_forced_token + _apply_ratio_restriction = _Real._apply_ratio_restriction + + +class TestGetForcedToken: + """Tests for the stateless _get_forced_token method.""" + + def setup_method(self): + self.model = FakeSamplerModel(is_comprehension=False) + + def test_no_trigger_returns_none(self): + assert self.model._get_forced_token([1, 2, 3]) is None + + def test_empty_history_returns_none(self): + assert self.model._get_forced_token([]) is None + + def test_end_of_think_forces_recaption(self): + assert self.model._get_forced_token([END_OF_THINK]) == RECAPTION + + def test_end_of_think_completed(self): + assert self.model._get_forced_token([END_OF_THINK, RECAPTION]) is None + + def test_end_of_recaption_forces_answer(self): + tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION] + assert self.model._get_forced_token(tokens) == ANSWER + + def test_end_of_recaption_forces_boi_after_answer(self): + tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION, ANSWER] + assert self.model._get_forced_token(tokens) == BOI + + def test_end_of_recaption_forces_size_after_boi(self): + tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION, ANSWER, BOI] + assert self.model._get_forced_token(tokens) == SIZE_TOKEN + + def test_full_sequence_complete(self): + tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION, ANSWER, BOI, SIZE_TOKEN] + assert self.model._get_forced_token(tokens) is None + + def test_diverged_history_returns_none(self): + tokens = [END_OF_RECAPTION, 999] # 999 != ANSWER + assert self.model._get_forced_token(tokens) is None + + def test_later_trigger_takes_precedence(self): + tokens = [END_OF_THINK, RECAPTION, END_OF_RECAPTION] + assert self.model._get_forced_token(tokens) == ANSWER + + def test_trigger_with_extra_tokens_before(self): + tokens = [1, 2, 3, END_OF_THINK] + assert self.model._get_forced_token(tokens) == RECAPTION + + +class TestComprehensionBlocking: + """Tests for comprehension mode token blocking.""" + + def test_blocked_tokens_masked(self): + model = FakeSamplerModel(is_comprehension=True) + vocab_size = 300 + logits = torch.zeros(1, vocab_size) + logits[0, BOI] = 5.0 + logits[0, SIZE_TOKEN] = 3.0 + logits[0, RATIO_START] = 2.0 + min_score = torch.finfo(logits.dtype).min + + for tid in model._blocked_token_ids: + if tid < vocab_size: + logits[0, tid] = min_score + + assert logits[0, BOI].item() == min_score + assert logits[0, SIZE_TOKEN].item() == min_score + assert logits[0, RATIO_START].item() == min_score + + def test_non_blocked_tokens_preserved(self): + model = FakeSamplerModel(is_comprehension=True) + vocab_size = 300 + logits = torch.zeros(1, vocab_size) + logits[0, 50] = 7.0 + min_score = torch.finfo(logits.dtype).min + + for tid in model._blocked_token_ids: + if tid < vocab_size: + logits[0, tid] = min_score + + assert logits[0, 50].item() == 7.0 + + +class TestRatioRestriction: + """Tests for _apply_ratio_restriction (greedy: only argmax ratio survives).""" + + def test_greedy_selects_single_ratio_token(self): + model = FakeSamplerModel(is_comprehension=False) + vocab_size = 300 + logits = torch.zeros(1, vocab_size) + logits[0, RATIO_START + 3] = 10.0 + logits[0, RATIO_START + 1] = 5.0 + logits[0, 50] = 20.0 # non-ratio, should be masked + min_score = torch.finfo(logits.dtype).min + + model._apply_ratio_restriction(logits, 0, min_score) + + assert logits[0, RATIO_START + 3].item() == 0 + assert logits[0, RATIO_START + 1].item() == min_score + assert logits[0, 50].item() == min_score + + def test_extra_ratio_slices_considered(self): + model = FakeSamplerModel(is_comprehension=False) + vocab_size = 300 + logits = torch.zeros(1, vocab_size) + logits[0, RATIO_OTHER_START] = 15.0 + logits[0, RATIO_START] = 5.0 + min_score = torch.finfo(logits.dtype).min + + model._apply_ratio_restriction(logits, 0, min_score) + + assert logits[0, RATIO_OTHER_START].item() == 0 + assert logits[0, RATIO_START].item() == min_score + + +class TestForceEosAfterRatio: + """Tests that a ratio token as last_token forces EOS.""" + + def test_ratio_token_forces_eos(self): + model = FakeSamplerModel(is_comprehension=False) + vocab_size = 300 + logits = torch.randn(1, vocab_size) + min_score = torch.finfo(logits.dtype).min + + logits[0].fill_(min_score) + logits[0, model._eos_token_id] = 0 + + assert logits[0, EOS].item() == 0 + non_eos_max = logits[0, :EOS].max().item() + assert non_eos_max == min_score diff --git a/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py b/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py index 6898763e4..ec4f4693d 100644 --- a/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py +++ b/tests/e2e/offline_inference/test_hunyuanimage3_text2img.py @@ -17,7 +17,7 @@ PROMPT = "A brown and white dog is running on the grass" MODEL_NAME = "tencent/HunyuanImage-3.0" LOCAL_CLIP_PATH = "openai/clip-vit-base-patch32" REPO_ROOT = Path(__file__).resolve().parents[3] -STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_moe.yaml" +STAGE_CONFIG_PATH = REPO_ROOT / "vllm_omni" / "model_executor" / "stage_configs" / "hunyuan_image3_t2i.yaml" pytestmark = [pytest.mark.advanced_model, pytest.mark.diffusion] diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index 5c280ddcf..6304eeab2 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -77,7 +77,9 @@ from vllm.multimodal.processing import ( from vllm.sequence import IntermediateTensors from vllm.transformers_utils.tokenizer import get_tokenizer from vllm.utils.tensor_schema import TensorSchema +from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler from vllm_omni.model_executor.models.hunyuan_image3.autoencoder_kl_3d import AutoencoderKLConv3D from vllm_omni.model_executor.models.hunyuan_image3.siglip2 import LightProjector, Siglip2VisionTransformer @@ -175,8 +177,11 @@ class HunyuanModel(HunYuanModel): return True return False + skipped_unexpected: set[str] = set() + for name, loaded_weight in weights: if contains_unexpected_keyword(name, unexpected_keywords): + skipped_unexpected.add(name) continue if "rotary_emb.inv_freq" in name: @@ -362,6 +367,17 @@ class HunyuanModel(HunYuanModel): weight_loader = getattr(param, "weight_loader", default_weight_loader) weight_loader(param, loaded_weight) loaded_params.add(name) + + if skipped_unexpected: + logger.warning_once( + "Skipped %d weights matching unexpected_keywords " + "(e.g. vae, vision_model, patch_embed, timestep_emb). " + "If upstream renamed components, these may be silently " + "lost. Skipped names: %s", + len(skipped_unexpected), + sorted(skipped_unexpected)[:10], + ) + return loaded_params @@ -1149,6 +1165,8 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo HunyuanImage3Inputs: TypeAlias = HunyuanImage3PixelInputs + prefer_model_sampler = True + packed_modules_mapping = { "qkv_proj": [ "q_proj", @@ -1199,6 +1217,10 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo else: self.lm_head = PPMissingLayer() + # --- AR-stage components --- + # These are needed for image encoding in the AR stage. + # If a future text-only stage is added, gate on vllm_config.model_config.model_stage. + # vae self.vae = AutoencoderKLConv3D.from_config(config.vae) self.patch_embed = UNetDown( @@ -1226,6 +1248,63 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo self._mrope_joint_img_sep_token_id = tokenizer.convert_tokens_to_ids("") self._mrope_max_num_patches = config.vit_processor.get("max_num_patches", 729) + # Special token IDs for logits processors (stage transitions). + # These mirror the official tokenization_hunyuan_image_3.py setup. + self._end_of_think_id = tokenizer.convert_tokens_to_ids("") + self._recaption_id = tokenizer.convert_tokens_to_ids("") + self._end_of_recaption_id = tokenizer.convert_tokens_to_ids("") + self._answer_id = tokenizer.convert_tokens_to_ids("") + self._end_of_answer_id = tokenizer.convert_tokens_to_ids("") + image_base_size = getattr(config, "image_base_size", 1024) + self._size_token_id = tokenizer.convert_tokens_to_ids(f"") + self._start_ratio_id = tokenizer.convert_tokens_to_ids("") + self._end_ratio_id = tokenizer.convert_tokens_to_ids("") + ratio_33 = tokenizer.convert_tokens_to_ids("") + ratio_36 = tokenizer.convert_tokens_to_ids("") + self._ratio_other_slices = [(ratio_33, ratio_36 + 1)] + # Build the full set of ratio token IDs for use as stop tokens. + self._all_ratio_ids = set(range(self._start_ratio_id, self._end_ratio_id + 1)) + for s, e in self._ratio_other_slices: + self._all_ratio_ids.update(range(s, e)) + + # Determine mode: comprehension (I2T/T2T) vs generation (IT2I/T2I). + engine_output_type = getattr(vllm_config.model_config, "engine_output_type", None) + self._is_comprehension = engine_output_type in (None, "text") + + # For comprehension mode, block image generation tokens but allow + # text structure tokens (, , etc.) so the model can + # follow its natural generation pattern. Stop tokens in YAML will + # terminate at or EOS. + self._blocked_token_ids: set[int] = set() + if self._is_comprehension: + self._blocked_token_ids.update( + [ + self._mrope_boi_token_id, # + self._mrope_eoi_token_id, # + self._size_token_id, # + ] + ) + self._blocked_token_ids.update(self._all_ratio_ids) + + # For generation mode, build stage transition map. + # Official logic: → [], + # → [, , ] + # After , restrict vocab to ratio tokens only. + # Stage-transition forced sequences, keyed by trigger token. + self._stage_transitions: dict[int, list[int]] = {} + if not self._is_comprehension: + self._stage_transitions[self._end_of_think_id] = [ + self._recaption_id, + ] + self._stage_transitions[self._end_of_recaption_id] = [ + self._answer_id, + self._mrope_boi_token_id, + self._size_token_id, + ] + + self._sampler: Sampler | None = None + self._eos_token_id: int = tokenizer.eos_token_id + self._replace_rotary_embeddings() def _replace_rotary_embeddings(self): @@ -1257,6 +1336,12 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo head_dim, rope_theta, ) + if replaced == 0: + raise RuntimeError( + "HunyuanImage3: _replace_rotary_embeddings replaced 0 layers. " + "The custom interleaved 2D mRoPE is not active — model outputs " + "will be incorrect. Check that model.layers[*].self_attn.rotary_emb exists." + ) def _parse_and_validate_image_input( self, @@ -1274,6 +1359,10 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo if vit_pixel_values is None or vae_pixel_values is None: return None + # Handle empty batch (e.g., during profiling with 0 images / T2T mode) + if vit_pixel_values.numel() == 0 or vae_pixel_values.numel() == 0: + return None + return HunyuanImage3PixelInputs( type="pixel_values", pixel_values={ @@ -1472,6 +1561,112 @@ class HunyuanImage3ForConditionalGeneration(nn.Module, SupportsMultiModal, Suppo logits = self.logits_processor(self.lm_head, hidden_states) return logits + # ------------------------------------------------------------------ + # Custom sampler — applies HunyuanImage3-specific logits processors + # before the standard sampling step. + # + # Comprehension (I2T / T2T): + # Block generation-specific special tokens so sampling can't + # accidentally produce , , ratio tokens, etc. + # + # Generation (IT2I / T2I think): + # 1. _StageTransitionLogitsProcessor — force token sequences at + # transition boundaries ( → , etc.) + # 2. _ConditionalSliceVocabLogitsProcessor — after , + # restrict vocab to ratio tokens only (greedy). + # ------------------------------------------------------------------ + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> SamplerOutput | None: + if logits is None or logits.numel() == 0: + return None + + if self._sampler is None: + self._sampler = Sampler() + + min_score = torch.finfo(logits.dtype).min + + assert logits.shape[0] == 1, f"HunyuanImage3 sampler requires max_num_seqs=1, got batch size {logits.shape[0]}" + + for req_idx in range(logits.shape[0]): + decoded_tokens: list[int] = ( + sampling_metadata.output_token_ids[req_idx] if req_idx < len(sampling_metadata.output_token_ids) else [] + ) + last_token = decoded_tokens[-1] if decoded_tokens else -1 + + if self._is_comprehension: + for tid in self._blocked_token_ids: + logits[req_idx, tid] = min_score + else: + forced = self._get_forced_token(decoded_tokens) + if forced is not None: + logits[req_idx].fill_(min_score) + logits[req_idx, forced] = 0 + elif last_token == self._size_token_id: + self._apply_ratio_restriction(logits, req_idx, min_score) + elif last_token in self._all_ratio_ids: + logits[req_idx].fill_(min_score) + logits[req_idx, self._eos_token_id] = 0 + + return self._sampler(logits=logits, sampling_metadata=sampling_metadata) + + def _get_forced_token(self, decoded_tokens: list[int]) -> int | None: + """Derive the next forced token from output history (stateless). + + Scans decoded_tokens backwards for the most recent trigger token, + then prefix-matches the forced sequence against what followed. + Returns the next token to force, or None if the sequence is complete + or history has diverged from the expected forced sequence. + """ + for i in range(len(decoded_tokens) - 1, -1, -1): + trigger = decoded_tokens[i] + if trigger not in self._stage_transitions: + continue + + forced_seq = self._stage_transitions[trigger] + emitted = decoded_tokens[i + 1 :] + + matched = 0 + for expected, actual in zip(forced_seq, emitted): + if actual != expected: + # History diverged from the expected forced sequence. + # Stop applying transition forcing for safety. + return None + matched += 1 + + if matched < len(forced_seq): + return forced_seq[matched] + return None + + return None + + def _apply_ratio_restriction( + self, + logits: torch.Tensor, + req_idx: int, + min_score: float, + ) -> None: + """Port of official _ConditionalSliceVocabLogitsProcessor.__call__. + + After the size token, only allow ratio tokens and pick greedily. + """ + original = logits[req_idx].clone() + logits[req_idx].fill_(min_score) + # Allow primary ratio range. + logits[req_idx, self._start_ratio_id : self._end_ratio_id + 1] = original[ + self._start_ratio_id : self._end_ratio_id + 1 + ] + # Allow extra ratio slices. + for s, e in self._ratio_other_slices: + logits[req_idx, s:e] = original[s:e] + # Force greedy: keep only the argmax. + max_id = logits[req_idx].argmax().item() + logits[req_idx].fill_(min_score) + logits[req_idx, max_id] = 0 + def make_empty_intermediate_tensors( self, batch_size: int, dtype: torch.dtype, device: torch.device ) -> IntermediateTensors: diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image3_i2t.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_i2t.yaml new file mode 100644 index 000000000..203b54f25 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/hunyuan_image3_i2t.yaml @@ -0,0 +1,44 @@ +# Stage config for HunyuanImage-3.0 Image-to-Text (I2T / image understanding). +# Single LLM stage: AR model reads image + text prompt, generates text output. + +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + process: true + devices: "0,1,2,3" + max_batch_size: 1 + requires_multimodal_data: true + engine_args: + model_stage: AR + max_num_seqs: 1 + model_arch: HunyuanImage3ForCausalMM + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.95 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 4 + pipeline_parallel_size: 1 + hf_overrides: + rope_parameters: + mrope_section: [0, 32, 32] + rope_type: default + is_comprehension: true + final_output: true + final_output_type: text + default_sampling_params: + temperature: 0.0 + top_p: 0.95 + top_k: 1024 + max_tokens: 2048 + stop_token_ids: [127957, 128026] # <|endoftext|>, + detokenize: True + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image3_it2i.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_it2i.yaml new file mode 100644 index 000000000..9f6adece0 --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/hunyuan_image3_it2i.yaml @@ -0,0 +1,78 @@ +# Stage config for HunyuanImage-3.0 Image+Text-to-Image (image editing). +# Stage 0: AR (HunyuanImage3ForConditionalGeneration) — reads (image, text), emits latent tokens +# Stage 1: Diffusion (HunyuanImage3Pipeline / DiT + VAE) — denoise + decode latents → image + +stage_args: + # Stage 0: AR Model + - stage_id: 0 + stage_type: llm + runtime: + process: true + devices: "0,1,2,3" + max_batch_size: 1 + requires_multimodal_data: true # AR needs the original image + engine_args: + model_stage: AR + model_arch: HunyuanImage3ForCausalMM + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.95 + enforce_eager: true + trust_remote_code: true + engine_output_type: latent # AR outputs latent for DiT + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 4 + pipeline_parallel_size: 1 + hf_overrides: + rope_parameters: + mrope_section: [0, 32, 32] + rope_type: default + is_comprehension: false # Generation task, not comprehension + final_output: false # AR is not the final output + default_sampling_params: + temperature: 0.6 + top_p: 0.95 + top_k: 1024 + max_tokens: 4096 + stop_token_ids: [127957] # <|endoftext|> + detokenize: false + + # Stage 1: Diffusion (DiT + VAE) + # Receives latents from AR stage, performs denoising + VAE decode + - stage_id: 1 + stage_type: diffusion + runtime: + process: true + devices: "4,5,6,7" + max_batch_size: 1 + requires_multimodal_data: true # May need condition images + engine_args: + model_stage: dit + model_arch: HunyuanImage3ForCausalMM + enforce_eager: true + trust_remote_code: true + distributed_executor_backend: "mp" + parallel_config: + tensor_parallel_size: 4 + enable_expert_parallel: true + omni_kv_config: + need_recv_cache: true + engine_input_source: [0] # Input from AR stage + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.hunyuan_image3.ar2diffusion + final_output: true + final_output_type: image + default_sampling_params: + num_inference_steps: 50 + guidance_scale: 2.5 + +# Top-level runtime config +runtime: + enabled: true + defaults: + window_size: -1 # Trigger downstream only after full upstream completion + max_inflight: 1 # Process serially within each stage + edges: + - from: 0 # AR → Diffusion + to: 1 + window_size: -1 diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe.yaml deleted file mode 100644 index 808b4619f..000000000 --- a/vllm_omni/model_executor/stage_configs/hunyuan_image3_moe.yaml +++ /dev/null @@ -1,81 +0,0 @@ -# Stage config for running Hunyuan-Image3.0 for multi-stage omni runtime. -# Stage 0: AR Model (vLLM implementation) - -# The following config has been verified on 8x L40S-48G GPU. -modes: - - mode: text-to-image - stages: [1] - - mode: image-to-text - stages: [0] -stage_args: - - stage_id: 0 - stage_type: llm # Use llm stage type for AR stages - runtime: - process: true # Run this stage in a separate process - devices: "0,1,2,3,4,5,6,7" # Visible devices for this stage (CUDA_VISIBLE_DEVICES/torch.cuda.set_device) - engine_args: - model_stage: AR - max_num_seqs: 1 - model_arch: HunyuanImage3ForCausalMM - worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker - scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.3 - enforce_eager: true # Now we only support eager mode - trust_remote_code: true - engine_output_type: latent - enable_prefix_caching: false - max_num_batched_tokens: 32768 - tensor_parallel_size: 8 - pipeline_parallel_size: 1 - hf_overrides: - rope_parameters: - mrope_section: [0, 32, 32] - rope_type: default - is_comprehension: true - final_output: true - final_output_type: text - default_sampling_params: - temperature: 0.0 - top_p: 1.0 - top_k: -1 - max_tokens: 2048 - seed: 42 - detokenize: True - repetition_penalty: 1.1 - - stage_id: 1 - stage_type: diffusion - runtime: - process: true - devices: "0,1,2,3,4,5,6,7" - max_batch_size: 1 - engine_args: - model_stage: diffusion - enforce_eager: true - distributed_executor_backend: "mp" - vae_use_slicing: false - vae_use_tiling: false - cache_backend: null - cache_config: null - enable_cache_dit_summary: false - parallel_config: - pipeline_parallel_size: 1 - data_parallel_size: 1 - tensor_parallel_size: 8 - enable_expert_parallel: false - sequence_parallel_size: 1 - ulysses_degree: 1 - ring_degree: 1 - cfg_parallel_size: 1 - vae_patch_parallel_size: 1 - use_hsdp: false - hsdp_shard_size: -1 - hsdp_replicate_size: 1 - final_output: true - final_output_type: image - -# Top-level runtime config (concise): default windows and stage edges -runtime: - enabled: true - defaults: - window_size: -1 # Simplified: trigger downstream only after full upstream completion - max_inflight: 1 # Simplified: process serially within each stage diff --git a/vllm_omni/model_executor/stage_configs/hunyuan_image3_t2t.yaml b/vllm_omni/model_executor/stage_configs/hunyuan_image3_t2t.yaml new file mode 100644 index 000000000..60da8e0bc --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/hunyuan_image3_t2t.yaml @@ -0,0 +1,45 @@ +# Stage config for HunyuanImage-3.0 Text-to-Text (T2T / pure text generation). +# Single LLM stage: AR model reads text prompt only, generates text output. +# Sampling params aligned with official generation_config.json. + +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + process: true + devices: "0,1,2,3" + max_batch_size: 1 + requires_multimodal_data: false + engine_args: + model_stage: AR + max_num_seqs: 1 + model_arch: HunyuanImage3ForCausalMM + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.95 + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + max_num_batched_tokens: 32768 + tensor_parallel_size: 4 + pipeline_parallel_size: 1 + hf_overrides: + rope_parameters: + mrope_section: [0, 32, 32] + rope_type: default + is_comprehension: true + final_output: true + final_output_type: text + default_sampling_params: + temperature: 0.0 + top_p: 0.95 + top_k: 1024 + max_tokens: 2048 + stop_token_ids: [127957, 128026] # <|endoftext|>, + detokenize: True + +runtime: + enabled: true + defaults: + window_size: -1 + max_inflight: 1 diff --git a/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py b/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py new file mode 100644 index 000000000..89a7a28f6 --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/hunyuan_image3.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""Stage input processor for HunyuanImage3: AR → Diffusion transition. + +In IT2I (image editing) mode: + - Stage 0 (AR) receives (image + edit instruction), generates CoT/latent tokens + - Stage 1 (DiT) receives the AR output + original image, denoises → edited image + +The ar2diffusion function bridges these two stages, following the same +signature pattern as glm_image.ar2diffusion. +""" + +from typing import Any + +import torch +from vllm.inputs import TextPrompt +from vllm.logger import init_logger + +from vllm_omni.inputs.data import OmniTokensPrompt + +logger = init_logger(__name__) + + +def ar2diffusion( + stage_list: list[Any], + engine_input_source: list[int], + prompt: OmniTokensPrompt | TextPrompt | list | None = None, + requires_multimodal_data: bool = False, +) -> list[dict[str, Any]]: + """Process AR stage outputs to create Diffusion stage inputs. + + Args: + stage_list: List of stage clients (set by orchestrator). + engine_input_source: List of source stage IDs (from YAML). + prompt: Original user prompt (may contain multimodal data). + requires_multimodal_data: Whether to forward multimodal data. + + Returns: + List of dicts, each consumable by the HunyuanImage3 diffusion pipeline. + """ + if not engine_input_source: + raise ValueError("engine_input_source cannot be empty") + + source_stage_id = engine_input_source[0] + if source_stage_id >= len(stage_list): + raise IndexError(f"Invalid source stage_id: {source_stage_id}") + + if stage_list[source_stage_id].engine_outputs is None: + raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") + + ar_outputs = stage_list[source_stage_id].engine_outputs + diffusion_inputs = [] + + # Normalize prompt to list + if not isinstance(prompt, list): + prompt = [prompt] if prompt is not None else [{}] + + for i, ar_output in enumerate(ar_outputs): + output = ar_output.outputs[0] + generated_token_ids = output.token_ids + generated_text = getattr(output, "text", "") or "" + + # Get original prompt info + original_prompt = prompt[i] if i < len(prompt) else {} + if isinstance(original_prompt, dict): + pass + elif hasattr(original_prompt, "_asdict"): + original_prompt = original_prompt._asdict() + elif hasattr(original_prompt, "__dict__"): + original_prompt = vars(original_prompt) + else: + original_prompt = {} + + height = original_prompt.get("height", 1024) + width = original_prompt.get("width", 1024) + text_prompt = original_prompt.get("prompt", "") + + logger.info( + "[ar2diffusion] Request %d: AR generated %d tokens, text length=%d, target size=%dx%d", + i, + len(generated_token_ids), + len(generated_text), + height, + width, + ) + + token_tensor = torch.tensor(generated_token_ids, dtype=torch.long) + + diffusion_input: dict[str, Any] = { + "prompt": text_prompt, + "height": height, + "width": width, + "extra": { + "ar_token_ids": token_tensor, + "ar_generated_text": generated_text, + }, + } + + # Forward multimodal data (original image for IT2I conditioning) + mm_data = original_prompt.get("multi_modal_data") + if mm_data: + pil_image = mm_data.get("image") + if pil_image is None: + images = mm_data.get("images") + if images: + pil_image = images[0] if isinstance(images, list) else images + if pil_image is not None: + diffusion_input["pil_image"] = pil_image + + # Forward multimodal output from AR (if any) + if hasattr(ar_output, "multimodal_output") and ar_output.multimodal_output: + mm_output = ar_output.multimodal_output + if isinstance(mm_output, dict): + diffusion_input["extra"]["ar_multimodal_output"] = mm_output + + # Forward sampling params + for key in ["seed", "num_inference_steps", "guidance_scale", "negative_prompt"]: + if key in original_prompt: + diffusion_input[key] = original_prompt[key] + + diffusion_inputs.append(diffusion_input) + + return diffusion_inputs diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index eafff821a..d4ab78f13 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -1,6 +1,8 @@ import sys +from functools import cached_property from aenum import extend_enum +from vllm.config import ModelConfig as _OriginalModelConfig from vllm.inputs import TokensPrompt as _OriginalTokensPrompt from vllm.model_executor.layers.rotary_embedding import ( MRotaryEmbedding as _OriginalMRotaryEmbedding, @@ -17,6 +19,56 @@ from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.layers.rotary_embedding import OmniMRotaryEmbedding from vllm_omni.request import OmniRequest +# ============================================================================= +# Patch ModelConfig.is_mm_prefix_lm to support omni-specific models +# ============================================================================= +# WHY: HunyuanImage-3.0 requires bidirectional attention for image tokens +# (cond_token_attn_type: "joint_full" in config.json). vLLM gates this on +# is_mm_prefix_lm, which checks an internal MM_PREFIX_LM_MODELS list that +# does not include "hunyuan_image_3_moe" (the upstream HF model_type). +# +# WHY NOT model-level: is_mm_prefix_lm is checked in vLLM core (scheduler, +# attention backend selection) before model code runs — no model-level hook. +# +# SCOPE: Only affects model_type in _OMNI_MM_PREFIX_LM_MODELS (currently +# just "hunyuan_image_3_moe"). All other models fall through to the +# original vLLM implementation unchanged. +# +# FRAGILITY: Relies on is_mm_prefix_lm being a cached_property on +# ModelConfig. The __dict__ access + __set_name__ dance works around a +# pydantic dataclass issue in vllm 0.19.0+. If vLLM changes +# is_mm_prefix_lm to a regular method or removes it, this will break. +# +# TODO: Upstream a configurable MM_PREFIX_LM_MODELS or a model_config flag +# so this patch can be removed. +_OMNI_MM_PREFIX_LM_MODELS = ("hunyuan_image_3_moe",) +# Access via __dict__ to avoid triggering cached_property.__get__ which fails +# with "Cannot use cached_property instance without calling __set_name__" in +# pydantic dataclasses (vllm 0.19.0+). +_cp = _OriginalModelConfig.__dict__["is_mm_prefix_lm"] +_original_is_mm_prefix_lm = _cp.func if hasattr(_cp, "func") else _cp.fget + + +def _patched_is_mm_prefix_lm(self): + if _original_is_mm_prefix_lm(self): + return True + model_type = getattr(self.hf_config, "model_type", "") + return model_type in _OMNI_MM_PREFIX_LM_MODELS + + +_patched_cp = cached_property(_patched_is_mm_prefix_lm) +_patched_cp.__set_name__(_OriginalModelConfig, "is_mm_prefix_lm") +_OriginalModelConfig.is_mm_prefix_lm = _patched_cp + +# Sanity check: verify the patch is active. If vLLM changes the descriptor +# type or __set_name__ semantics, this will fail loudly at import time +# rather than silently falling back to unpatched behavior. +_installed = _OriginalModelConfig.__dict__.get("is_mm_prefix_lm") +assert _installed is _patched_cp, ( + "is_mm_prefix_lm patch failed to install — bidirectional attention " + "for HunyuanImage3 will not work. Check vLLM ModelConfig changes." +) + # ============================================================================= # Patch GlmImageTextConfig to expose mrope_section in rope_parameters # ============================================================================= From 817e32d548de74d374b34b6f7dcdccb8342cf4cd Mon Sep 17 00:00:00 2001 From: Zhang Jian Date: Thu, 16 Apr 2026 23:33:40 +0800 Subject: [PATCH 31/97] [Quantization] feat: add FP8 for Omnigen2 (#2441) Signed-off-by: Zhang --- .../models/omnigen2/omnigen2_transformer.py | 177 +++++++++++++++++- .../models/omnigen2/pipeline_omnigen2.py | 7 +- 2 files changed, 174 insertions(+), 10 deletions(-) diff --git a/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py b/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py index 9ff681a3c..3f03563a1 100644 --- a/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py +++ b/vllm_omni/diffusion/models/omnigen2/omnigen2_transformer.py @@ -5,6 +5,8 @@ from types import SimpleNamespace import torch import torch.nn as nn +import torch.nn.functional as F +import vllm._custom_ops as ops from diffusers.models.activations import get_activation from diffusers.models.embeddings import Timesteps, get_1d_rotary_pos_embed from diffusers.models.modeling_outputs import Transformer2DModelOutput @@ -16,6 +18,7 @@ from vllm.model_executor.layers.linear import ( QKVParallelLinear, RowParallelLinear, ) +from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm_omni.diffusion.attention.layer import Attention @@ -24,6 +27,105 @@ from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) +def _patch_cutlass_padded_fp8(): + """Monkey-patch vllm._custom_ops.cutlass_scaled_mm to pad tensors whose + dimensions are not multiples of 16, so the CUTLASS FP8 kernel is used. + + OmniGen2 has hidden_size=2520 (2520 % 16 == 8). Without this patch, + vLLM's cutlass_scaled_mm falls back to a Triton scaled_mm kernel for + every FP8 linear layer (QKV, attn output, gate_up_proj, down_proj), + which is dramatically slower than the native CUTLASS FP8 tensor-core + path on H100/H200 GPUs. + + Weight tensors (b) are constant across forward passes, so padded + versions are computed once and cached by data_ptr to avoid repeated + allocation and column-major conversion overhead. + """ + _orig_cutlass_scaled_mm = ops.cutlass_scaled_mm + # Cache: data_ptr → (padded_b, padded_bias, padded_scale_b, pad_k, pad_n, orig_n) + _weight_cache: dict[int, tuple] = {} + + def _padded_cutlass_scaled_mm( + a: torch.Tensor, + b: torch.Tensor, + scale_a: torch.Tensor, + scale_b: torch.Tensor, + out_dtype: torch.dtype, + bias: torch.Tensor | None = None, + ) -> torch.Tensor: + if b.shape[0] % 16 == 0 and b.shape[1] % 16 == 0: + return _orig_cutlass_scaled_mm(a, b, scale_a, scale_b, out_dtype, bias) + + # Reshape to 2D (mirrors the original function) + target_shape = (*a.shape[:-1], b.shape[1]) + a = a.view(-1, a.shape[-1]) + orig_n = b.shape[1] + + # Cache the padded weight — it's a model parameter that never changes. + key = b.data_ptr() + if key not in _weight_cache: + pad_k = (16 - b.shape[0] % 16) % 16 + pad_n = (16 - orig_n % 16) % 16 + b_pad = b + if pad_k > 0: + b_pad = F.pad(b_pad, (0, 0, 0, pad_k)) + if pad_n > 0: + b_pad = F.pad(b_pad, (0, pad_n)) + # CUTLASS requires b column-major (stride(0)==1). + b_pad = b_pad.t().contiguous().t() + + bias_pad = None + if bias is not None and pad_n > 0: + bias_pad = F.pad(bias, (0, pad_n)) + + scale_b_pad = scale_b + if scale_b.numel() > 1 and pad_n > 0: + scale_b_pad = F.pad( + scale_b.view(-1, scale_b.shape[-1]), + (0, pad_n), + value=1.0, + ) + + _weight_cache[key] = ( + b_pad, + bias_pad, + scale_b_pad, + pad_k, + pad_n, + orig_n, + ) + + b_pad, bias_pad, scale_b_pad, pad_k, pad_n, orig_n = _weight_cache[key] + + # Pad activations on K dimension (cheap — activations are small). + if pad_k > 0: + a = F.pad(a, (0, pad_k)).contiguous() + + out = torch.empty((a.shape[0], b_pad.shape[1]), dtype=out_dtype, device=a.device) + torch.ops._C.cutlass_scaled_mm( + out, + a, + b_pad, + scale_a, + scale_b_pad, + bias_pad if bias is not None else None, + ) + + if pad_n > 0: + out = out[:, :orig_n] + + return out.view(*target_shape) + + ops.cutlass_scaled_mm = _padded_cutlass_scaled_mm + logger.info( + "Patched vllm._custom_ops.cutlass_scaled_mm with CUTLASS-padded FP8 " + "variant (avoids slow Triton fallback for non-%%16 dimensions)" + ) + + +_patch_cutlass_padded_fp8() + + class OmniGen2Attention(nn.Module): def __init__( self, @@ -31,6 +133,8 @@ class OmniGen2Attention(nn.Module): num_heads: int, num_kv_heads: int, eps: float = 1e-5, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ): super().__init__() self.dim = dim @@ -46,12 +150,26 @@ class OmniGen2Attention(nn.Module): total_num_kv_heads=num_kv_heads, disable_tp=True, bias=False, + quant_config=quant_config, + prefix=f"{prefix}.to_qkv", ) self.norm_q = RMSNorm(self.head_dim, eps=eps) self.norm_k = RMSNorm(self.head_dim, eps=eps) - self.to_out = nn.ModuleList([nn.Linear(dim, dim, bias=False)]) + self.to_out = nn.ModuleList( + [ + RowParallelLinear( + dim, + dim, + bias=False, + input_is_parallel=False, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.to_out.0", + ) + ] + ) self.attn = Attention( num_heads=num_heads, head_size=self.head_dim, @@ -78,6 +196,9 @@ class OmniGen2Attention(nn.Module): """ batch_size = hidden_states.shape[0] + # Contiguous layout for FP8 quantized linear GEMMs (matches FLUX DiT). + hidden_states = hidden_states.contiguous() + # Get Query-Key-Value Pair qkv, _ = self.to_qkv(hidden_states) @@ -121,7 +242,7 @@ class OmniGen2Attention(nn.Module): hidden_states = hidden_states.reshape(batch_size, -1, self.num_heads * self.head_dim) hidden_states = hidden_states.to(dtype) - hidden_states = self.to_out[0](hidden_states) + hidden_states = self.to_out[0](hidden_states.contiguous()) return hidden_states @@ -233,6 +354,7 @@ class LuminaRMSNormZero(nn.Module): embedding_dim: int, norm_eps: float, norm_elementwise_affine: bool, + **kwargs, ): super().__init__() self.silu = nn.SiLU() @@ -325,6 +447,8 @@ class LuminaFeedForward(nn.Module): inner_dim: int, multiple_of: int | None = 256, ffn_dim_multiplier: float | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ): super().__init__() @@ -338,6 +462,8 @@ class LuminaFeedForward(nn.Module): [inner_dim, inner_dim], bias=False, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.gate_up_proj", ) self.act_fn = get_act_and_mul_fn("silu") self.down_proj = RowParallelLinear( @@ -346,6 +472,8 @@ class LuminaFeedForward(nn.Module): bias=False, input_is_parallel=True, return_bias=False, + quant_config=quant_config, + prefix=f"{prefix}.down_proj", ) def forward(self, x): @@ -591,6 +719,8 @@ class OmniGen2TransformerBlock(nn.Module): ffn_dim_multiplier: float, norm_eps: float, modulation: bool = True, + quant_config: QuantizationConfig | None = None, + prefix: str = "", ) -> None: """Initialize the transformer block.""" super().__init__() @@ -602,6 +732,8 @@ class OmniGen2TransformerBlock(nn.Module): num_heads=num_attention_heads, num_kv_heads=num_kv_heads, eps=1e-5, + quant_config=quant_config, + prefix=f"{prefix}.attn", ) # Initialize feed-forward network @@ -610,11 +742,19 @@ class OmniGen2TransformerBlock(nn.Module): inner_dim=4 * dim, multiple_of=multiple_of, ffn_dim_multiplier=ffn_dim_multiplier, + quant_config=quant_config, + prefix=f"{prefix}.feed_forward", ) # Initialize normalization layers if modulation: - self.norm1 = LuminaRMSNormZero(embedding_dim=dim, norm_eps=norm_eps, norm_elementwise_affine=True) + self.norm1 = LuminaRMSNormZero( + embedding_dim=dim, + norm_eps=norm_eps, + norm_elementwise_affine=True, + quant_config=quant_config, + prefix=f"{prefix}.norm1", + ) else: self.norm1 = RMSNorm(dim, eps=norm_eps) @@ -713,6 +853,7 @@ class OmniGen2Transformer2DModel(nn.Module): axes_lens: tuple[int, int, int] = (1024, 1664, 1664), text_feat_dim: int = 2048, timestep_scale: float = 1000.0, + quant_config: QuantizationConfig | None = None, ) -> None: """Initialize the OmniGen2 transformer model.""" super().__init__() @@ -770,8 +911,10 @@ class OmniGen2Transformer2DModel(nn.Module): ffn_dim_multiplier, norm_eps, modulation=True, + quant_config=quant_config, + prefix=f"noise_refiner.{i}", ) - for _ in range(num_refiner_layers) + for i in range(num_refiner_layers) ] ) @@ -785,8 +928,10 @@ class OmniGen2Transformer2DModel(nn.Module): ffn_dim_multiplier, norm_eps, modulation=True, + quant_config=quant_config, + prefix=f"ref_image_refiner.{i}", ) - for _ in range(num_refiner_layers) + for i in range(num_refiner_layers) ] ) @@ -800,8 +945,10 @@ class OmniGen2Transformer2DModel(nn.Module): ffn_dim_multiplier, norm_eps, modulation=False, + quant_config=quant_config, + prefix=f"context_refiner.{i}", ) - for _ in range(num_refiner_layers) + for i in range(num_refiner_layers) ] ) @@ -816,8 +963,10 @@ class OmniGen2Transformer2DModel(nn.Module): ffn_dim_multiplier, norm_eps, modulation=True, + quant_config=quant_config, + prefix=f"layers.{i}", ) - for _ in range(num_layers) + for i in range(num_layers) ] ) @@ -847,11 +996,25 @@ class OmniGen2Transformer2DModel(nn.Module): temb, ): batch_size = len(hidden_states) + has_ref_tokens = any(ref_img_len > 0 for ref_lens in l_effective_ref_img_len for ref_img_len in ref_lens) max_combined_img_len = max( [img_len + sum(ref_img_len) for img_len, ref_img_len in zip(l_effective_img_len, l_effective_ref_img_len)] ) hidden_states = self.x_embedder(hidden_states) + if not has_ref_tokens: + # FP8 kernels do not support zero-token GEMM on ref_image_patch_embedder; skip that path only. + # Still run noise_refiner and return the same combined layout as the no-ref case below + # (batch, max_combined_img_len, hidden) — not raw noise tokens alone. + for layer in self.noise_refiner: + hidden_states = layer(hidden_states, padded_img_mask, noise_rotary_emb, temb) + combined_img_hidden_states = hidden_states.new_zeros( + batch_size, max_combined_img_len, self.config.hidden_size + ) + for i, img_len in enumerate(l_effective_img_len): + combined_img_hidden_states[i, :img_len] = hidden_states[i, :img_len] + return combined_img_hidden_states + ref_image_hidden_states = self.ref_image_patch_embedder(ref_image_hidden_states) for i in range(batch_size): diff --git a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py index e8e307b87..04720c932 100644 --- a/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py +++ b/vllm_omni/diffusion/models/omnigen2/pipeline_omnigen2.py @@ -676,7 +676,10 @@ class OmniGen2Pipeline(CFGParallelMixin, nn.Module): ) transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, OmniGen2Transformer2DModel) - self.transformer = OmniGen2Transformer2DModel(**transformer_kwargs) + self.transformer = OmniGen2Transformer2DModel( + **transformer_kwargs, + quant_config=od_config.quantization_config, + ) self.mllm = Qwen2_5_VLForConditionalGeneration.from_pretrained( model, subfolder="mllm", local_files_only=local_files_only ).to(self.device) @@ -1253,8 +1256,6 @@ class OmniGen2Pipeline(CFGParallelMixin, nn.Module): # broadcast to batch dimension in a way that's compatible with ONNX/Core ML timestep = t.expand(latents.shape[0]).to(latents.dtype) - batch_size, num_channels_latents, height, width = latents.shape - optional_kwargs = {} if "ref_image_hidden_states" in set(inspect.signature(self.transformer.forward).parameters.keys()): optional_kwargs["ref_image_hidden_states"] = ref_image_hidden_states From bf64a6b841d5ddb705ac1d6f0c48b484a333ccd9 Mon Sep 17 00:00:00 2001 From: Lancer Date: Fri, 17 Apr 2026 07:55:00 +0800 Subject: [PATCH 32/97] [Feature] Flux2 klein inpaint (#1180) Signed-off-by: Lancer Co-authored-by: Hongsheng Liu --- docs/serving/image_edit_api.md | 2 + .../test_flux2_klein_inpaint.py | 216 ++++++++++++++++++ .../test_flux2_klein_inpaint_expansion.py | 160 +++++++++++++ .../flux2_klein/pipeline_flux2_klein.py | 199 +++++++++++++++- vllm_omni/entrypoints/openai/api_server.py | 11 + vllm_omni/entrypoints/openai/serving_chat.py | 2 +- 6 files changed, 583 insertions(+), 7 deletions(-) create mode 100644 tests/e2e/offline_inference/test_flux2_klein_inpaint.py create mode 100644 tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py diff --git a/docs/serving/image_edit_api.md b/docs/serving/image_edit_api.md index d254ac06a..79303e1a6 100644 --- a/docs/serving/image_edit_api.md +++ b/docs/serving/image_edit_api.md @@ -104,6 +104,8 @@ Content-Type: multipart/form-data | `guidance_scale` | float | model defaults | Classifier-free guidance scale (typically 0.0-20.0) | | `true_cfg_scale` | float | model defaults | True CFG scale (model-specific parameter, may be ignored if not supported) | | `seed` | integer | null | Random seed for reproducibility | +| `reference_image` | string or array | null | Reference image for inpainting | +| `mask_image` | string or array | null | Mask for inpainting (white areas will be inpainted) | ### Response Format diff --git a/tests/e2e/offline_inference/test_flux2_klein_inpaint.py b/tests/e2e/offline_inference/test_flux2_klein_inpaint.py new file mode 100644 index 000000000..ac7169361 --- /dev/null +++ b/tests/e2e/offline_inference/test_flux2_klein_inpaint.py @@ -0,0 +1,216 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end test for Flux2 Klein inpainting. + +""" + +# ruff: noqa: E402 + +import os +import sys +from pathlib import Path + +import pytest +import torch +from PIL import Image, ImageDraw + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams +from vllm_omni.outputs import OmniRequestOutput +from vllm_omni.platforms import current_omni_platform + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from vllm_omni import Omni + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +MODEL = "black-forest-labs/FLUX.2-klein-4B" + +_HEIGHT = 512 +_WIDTH = 512 +_NUM_INFERENCE_STEPS = 4 + + +def _create_test_image(width: int = _WIDTH, height: int = _HEIGHT, color: tuple = (128, 128, 128)) -> Image.Image: + return Image.new("RGB", (width, height), color) + + +def _create_test_mask(width: int = _WIDTH, height: int = _HEIGHT) -> Image.Image: + mask = Image.new("L", (width, height), 0) + draw = ImageDraw.Draw(mask) + draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255) + return mask + + +def _create_test_inputs(color: tuple = (100, 150, 200)): + return _create_test_image(_WIDTH, _HEIGHT, color), _create_test_mask(_WIDTH, _HEIGHT) + + +def _extract_images_from_output(outputs: list) -> list[Image.Image]: + images = [] + for req_output in outputs: + if hasattr(req_output, "images") and req_output.images: + images.extend(req_output.images) + elif hasattr(req_output, "request_output") and req_output.request_output: + stage_out = req_output.request_output + if isinstance(stage_out, OmniRequestOutput) and hasattr(stage_out, "images"): + images.extend(stage_out.images) + elif isinstance(stage_out, list): + for s in stage_out: + if hasattr(s, "images") and s.images: + images.extend(s.images) + return images + + +@pytest.mark.core_model +@pytest.mark.diffusion +def test_flux2_klein_inpaint_basic(): + m = None + try: + m = Omni(model=MODEL) + input_image, mask_image = _create_test_inputs() + + outputs = m.generate( + prompts=[ + { + "prompt": "Fill in the masked area with a beautiful garden", + "multi_modal_data": {"image": input_image, "mask_image": mask_image}, + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=0.0, + generator=torch.Generator(current_omni_platform.device_type).manual_seed(42), + num_outputs_per_prompt=1, + ), + ) + + images = _extract_images_from_output(list(outputs)) + assert len(images) == 1 + assert images[0].size == (_WIDTH, _HEIGHT) + finally: + if m is not None and hasattr(m, "close"): + m.close() + + +@pytest.mark.diffusion +def test_flux2_klein_inpaint_deterministic(): + m = None + try: + m = Omni(model=MODEL) + input_image, mask_image = _create_test_inputs() + seed = 12345 + + gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(seed) + gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(seed) + + outputs1 = m.generate( + prompts=[ + { + "prompt": "A red flower in a field", + "multi_modal_data": {"image": input_image, "mask_image": mask_image}, + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=0.0, + generator=gen1, + num_outputs_per_prompt=1, + ), + ) + + outputs2 = m.generate( + prompts=[ + { + "prompt": "A red flower in a field", + "multi_modal_data": {"image": input_image, "mask_image": mask_image}, + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=0.0, + generator=gen2, + num_outputs_per_prompt=1, + ), + ) + + images1 = _extract_images_from_output(list(outputs1)) + images2 = _extract_images_from_output(list(outputs2)) + + assert len(images1) == 1 + assert len(images2) == 1 + + assert list(images1[0].getdata()) == list(images2[0].getdata()), ( + "Same input with same seed should produce identical output. " + "This is critical for offline/online consistency." + ) + finally: + if m is not None and hasattr(m, "close"): + m.close() + + +@pytest.mark.diffusion +def test_flux2_klein_inpaint_different_seeds_different_output(): + m = None + try: + m = Omni(model=MODEL) + input_image, mask_image = _create_test_inputs() + + gen1 = torch.Generator(current_omni_platform.device_type).manual_seed(42) + gen2 = torch.Generator(current_omni_platform.device_type).manual_seed(99999) + + outputs1 = m.generate( + prompts=[ + { + "prompt": "A beautiful landscape", + "multi_modal_data": {"image": input_image, "mask_image": mask_image}, + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=0.0, + generator=gen1, + num_outputs_per_prompt=1, + ), + ) + + outputs2 = m.generate( + prompts=[ + { + "prompt": "A beautiful landscape", + "multi_modal_data": {"image": input_image, "mask_image": mask_image}, + } + ], + sampling_params_list=OmniDiffusionSamplingParams( + height=_HEIGHT, + width=_WIDTH, + num_inference_steps=_NUM_INFERENCE_STEPS, + guidance_scale=0.0, + generator=gen2, + num_outputs_per_prompt=1, + ), + ) + + images1 = _extract_images_from_output(list(outputs1)) + images2 = _extract_images_from_output(list(outputs2)) + + assert len(images1) == 1 + assert len(images2) == 1 + + different_pixel_count = sum(1 for p1, p2 in zip(images1[0].getdata(), images2[0].getdata()) if p1 != p2) + assert different_pixel_count > 0, "Different seeds should produce different outputs" + finally: + if m is not None and hasattr(m, "close"): + m.close() diff --git a/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py b/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py new file mode 100644 index 000000000..f59a0e783 --- /dev/null +++ b/tests/e2e/online_serving/test_flux2_klein_inpaint_expansion.py @@ -0,0 +1,160 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +""" +End-to-end tests for Flux2 Klein inpainting in online serving mode. + +Uses /v1/images/edits endpoint which is the correct API for image inpainting. +""" + +import base64 +from io import BytesIO + +import httpx +import pytest +from PIL import Image, ImageDraw + +from tests.conftest import OmniServer, OmniServerParams + +MODEL = "black-forest-labs/FLUX.2-klein-4B" + +_HEIGHT = 512 +_WIDTH = 512 +_NUM_INFERENCE_STEPS = 4 + + +def _get_diffusion_feature_cases(model: str): + return [ + pytest.param( + OmniServerParams( + model=model, + server_args=["--tensor-parallel-size", "2"], + ), + id="tp2_basic", + ), + ] + + +def _image_to_base64_jpeg(image: Image.Image) -> str: + buffer = BytesIO() + image.save(buffer, format="JPEG") + buffer.seek(0) + return base64.b64encode(buffer.read()).decode("utf-8") + + +def _create_test_mask_base64(width: int = _WIDTH, height: int = _HEIGHT) -> str: + mask = Image.new("L", (width, height), 0) + draw = ImageDraw.Draw(mask) + draw.rectangle([width // 4, height // 4, width * 3 // 4, height * 3 // 4], fill=255) + return _image_to_base64_jpeg(mask) + + +def _compare_images(img1: Image.Image, img2: Image.Image) -> bool: + return list(img1.getdata()) == list(img2.getdata()) + + +def _send_edit_request(host: str, port: int, model: str, image_b64: str, mask_b64: str, prompt: str, **kwargs): + url = f"http://{host}:{port}/v1/images/edits" + files = { + "image": ("image.jpg", base64.b64decode(image_b64), "image/jpeg"), + "mask_image": ("mask.jpg", base64.b64decode(mask_b64), "image/jpeg"), + } + data = {"prompt": prompt, "model": model, **kwargs} + with httpx.Client(timeout=60.0) as client: + response = client.post(url, files=files, data=data) + response.raise_for_status() + return response.json() + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) +def test_flux2_klein_inpaint_basic(omni_server: OmniServer): + input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) + mask_b64 = _create_test_mask_base64() + + result = _send_edit_request( + host=omni_server.host, + port=omni_server.port, + model=MODEL, + image_b64=input_image_b64, + mask_b64=mask_b64, + prompt="Fill in the masked area with a beautiful garden", + guidance_scale=1.0, + num_inference_steps=_NUM_INFERENCE_STEPS, + n=1, + seed=42, + ) + + assert "data" in result and len(result["data"]) == 1 + img_data = result["data"][0].get("b64_json") or result["data"][0].get("url", "").split(",")[-1] + img = Image.open(BytesIO(base64.b64decode(img_data))) + assert img.size == (_WIDTH, _HEIGHT) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) +def test_flux2_klein_inpaint_deterministic(omni_server: OmniServer): + input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) + mask_b64 = _create_test_mask_base64() + prompt = "A red flower in a field" + + result1 = _send_edit_request( + host=omni_server.host, + port=omni_server.port, + model=MODEL, + image_b64=input_image_b64, + mask_b64=mask_b64, + prompt=prompt, + guidance_scale=1.0, + num_inference_steps=_NUM_INFERENCE_STEPS, + n=1, + seed=12345, + ) + + result2 = _send_edit_request( + host=omni_server.host, + port=omni_server.port, + model=MODEL, + image_b64=input_image_b64, + mask_b64=mask_b64, + prompt=prompt, + guidance_scale=1.0, + num_inference_steps=_NUM_INFERENCE_STEPS, + n=1, + seed=12345, + ) + + img1_data = result1["data"][0].get("b64_json") or result1["data"][0].get("url", "").split(",")[-1] + img2_data = result2["data"][0].get("b64_json") or result2["data"][0].get("url", "").split(",")[-1] + + img1 = Image.open(BytesIO(base64.b64decode(img1_data))) + img2 = Image.open(BytesIO(base64.b64decode(img2_data))) + + assert _compare_images(img1, img2), ( + "Same input with same seed should produce identical output. This is critical for offline/online consistency." + ) + + +@pytest.mark.advanced_model +@pytest.mark.diffusion +@pytest.mark.parametrize("omni_server", _get_diffusion_feature_cases(MODEL), indirect=True) +def test_flux2_klein_inpaint_multiple_outputs(omni_server: OmniServer): + input_image_b64 = _image_to_base64_jpeg(Image.new("RGB", (_WIDTH, _HEIGHT), (128, 128, 128))) + mask_b64 = _create_test_mask_base64() + + result = _send_edit_request( + host=omni_server.host, + port=omni_server.port, + model=MODEL, + image_b64=input_image_b64, + mask_b64=mask_b64, + prompt="A beautiful landscape", + guidance_scale=1.0, + num_inference_steps=_NUM_INFERENCE_STEPS, + n=2, + seed=42, + ) + + assert "data" in result and len(result["data"]) == 2 diff --git a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py index 437dd58d0..021d28c2a 100644 --- a/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py +++ b/vllm_omni/diffusion/models/flux2_klein/pipeline_flux2_klein.py @@ -234,7 +234,15 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi self.transformer = Flux2Transformer2DModel(quant_config=od_config.quantization_config, **transformer_kwargs) self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) if getattr(self, "vae", None) else 8 + self.latent_channels = self.vae.config.latent_channels if hasattr(self.vae, "config") else 16 self.image_processor = Flux2ImageProcessor(vae_scale_factor=self.vae_scale_factor * 2) + self.mask_processor = VaeImageProcessor( + vae_scale_factor=self.vae_scale_factor * 2, + vae_latent_channels=self.latent_channels, + do_normalize=False, + do_binarize=True, + do_convert_grayscale=True, + ) self.tokenizer_max_length = 512 self.default_sample_size = 128 @@ -247,6 +255,14 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler ) + def get_timesteps(self, num_inference_steps, strength, device): + init_timestep = min(num_inference_steps * strength, num_inference_steps) + t_start = int(max(num_inference_steps - init_timestep, 0)) + timesteps = self.scheduler.timesteps[t_start * self.scheduler.order :] + if hasattr(self.scheduler, "set_begin_index"): + self.scheduler.set_begin_index(t_start * self.scheduler.order) + return timesteps, num_inference_steps - t_start + @staticmethod def _get_qwen3_prompt_embeds( text_encoder: Qwen3ForCausalLM, @@ -582,6 +598,54 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi return image_latents, image_latent_ids + def prepare_mask_latents( + self, + mask, + masked_image, + batch_size, + num_channels_latents, + num_images_per_prompt, + height, + width, + dtype, + device, + generator, + ): + height = 2 * (int(height) // (self.vae_scale_factor * 2)) + width = 2 * (int(width) // (self.vae_scale_factor * 2)) + mask = torch.nn.functional.interpolate(mask, size=(height, width)) + mask = mask.to(device=device, dtype=dtype) + + batch_size = batch_size * num_images_per_prompt + + if masked_image is not None: + masked_image = masked_image.to(device=device, dtype=dtype) + if masked_image.shape[1] != num_channels_latents: + masked_image_latents = self._encode_vae_image(image=masked_image, generator=generator) + else: + masked_image_latents = masked_image + else: + masked_image_latents = None + + if mask.shape[0] < batch_size: + if not batch_size % mask.shape[0] == 0: + raise ValueError("The passed mask and the required batch size don't match.") + mask = mask.repeat(batch_size // mask.shape[0], 1, 1, 1) + if masked_image_latents is not None and masked_image_latents.shape[0] < batch_size: + if not batch_size % masked_image_latents.shape[0] == 0: + raise ValueError("The passed mask and the required batch size don't match.") + masked_image_latents = masked_image_latents.repeat(batch_size // masked_image_latents.shape[0], 1, 1, 1) + + if masked_image_latents is not None: + masked_image_latents = masked_image_latents.to(device=device, dtype=dtype) + masked_image_latents = self._pack_latents(masked_image_latents) + + mask = mask.repeat(1, self.latent_channels, 1, 1) + mask = self._patchify_latents(mask) + mask = self._pack_latents(mask) + + return mask, masked_image_latents + def check_inputs( self, prompt, @@ -590,7 +654,15 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi prompt_embeds=None, callback_on_step_end_tensor_inputs=None, guidance_scale=None, + strength=None, + num_inference_steps=None, ): + if strength is not None: + if strength < 0 or strength > 1: + raise ValueError(f"strength must be between 0 and 1, got {strength}") + if num_inference_steps is not None and num_inference_steps <= 0: + raise ValueError(f"num_inference_steps must be positive, got {num_inference_steps}") + if ( height is not None and height % (self.vae_scale_factor * 2) != 0 @@ -622,7 +694,7 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi elif prompt is not None and (not isinstance(prompt, str) and not isinstance(prompt, list)): raise ValueError(f"`prompt` has to be of type `str` or `list` but is {type(prompt)}") - if guidance_scale > 1.0 and self.is_distilled: + if guidance_scale is not None and guidance_scale > 1.0 and self.is_distilled: logger.warning(f"Guidance scale {guidance_scale} is ignored for step-wise distilled models.") @property @@ -653,11 +725,14 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi self, req: OmniDiffusionRequest, image: PIL.Image.Image | list[PIL.Image.Image] | None = None, + reference_image: PIL.Image.Image | list[PIL.Image.Image] | None = None, + mask_image: PIL.Image.Image | list[PIL.Image.Image] | None = None, prompt: str | list[str] | None = None, height: int | None = None, width: int | None = None, num_inference_steps: int = 50, sigmas: list[float] | None = None, + strength: float = 1.0, guidance_scale: float | None = 4.0, num_images_per_prompt: int = 1, generator: torch.Generator | list[torch.Generator] | None = None, @@ -671,6 +746,7 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi callback_on_step_end_tensor_inputs: list[str] = ["latents"], max_sequence_length: int = 512, text_encoder_out_layers: tuple[int, ...] = (9, 18, 27), + padding_mask_crop: int | None = None, ) -> DiffusionOutput: r""" Function invoked when calling the pipeline for generation. @@ -804,6 +880,8 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi prompt_embeds=prompt_embeds, callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs, guidance_scale=guidance_scale, + strength=strength, + num_inference_steps=num_inference_steps, ) self._guidance_scale = guidance_scale @@ -848,6 +926,9 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi if image is not None and not isinstance(image, list): image = [image] + multiple_of = self.vae_scale_factor * 2 + crops_coords = None + resize_mode = "crop" condition_images = None if image is not None: for img in image: @@ -860,10 +941,14 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi img = self.image_processor._resize_to_target_area(img, 1024 * 1024) image_width, image_height = img.size - multiple_of = self.vae_scale_factor * 2 image_width = (image_width // multiple_of) * multiple_of image_height = (image_height // multiple_of) * multiple_of - img = self.image_processor.preprocess(img, height=image_height, width=image_width, resize_mode="crop") + if padding_mask_crop is not None: + crops_coords = self.mask_processor.get_crop_region(mask_image, width, height, pad=padding_mask_crop) + resize_mode = "fill" + img = self.image_processor.preprocess( + img, height=image_height, width=image_width, crops_coords=crops_coords, resize_mode=resize_mode + ) condition_images.append(img) height = height or image_height width = width or image_width @@ -871,6 +956,16 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi height = height or self.default_sample_size * self.vae_scale_factor width = width or self.default_sample_size * self.vae_scale_factor + # Get mask_image and reference_image + multi_modal_data = req.prompts[0].get("multi_modal_data", {}) if req.prompts else {} + mask_image = multi_modal_data.get("mask_image") + reference_image = multi_modal_data.get("reference_image") + + if mask_image is not None and (image is None or (isinstance(image, list) and len(image) == 0)): + raise ValueError("image must be provided when using mask_image for inpainting") + + init_image = condition_images[0] if condition_images else None + # 5. prepare latent variables num_channels_latents = self.transformer.config.in_channels // 4 latents, latent_ids = self.prepare_latents( @@ -884,6 +979,8 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi latents=latents, ) + original_latent_ids = latent_ids + image_latents = None image_latent_ids = None if condition_images is not None: @@ -895,6 +992,71 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi dtype=self.vae.dtype, ) + # Preprocess reference_image + if reference_image is not None and not ( + isinstance(reference_image, torch.Tensor) and reference_image.size(1) == self.latent_channels + ): + if ( + isinstance(reference_image, list) + and isinstance(reference_image[0], torch.Tensor) + and reference_image[0].ndim == 4 + ): + reference_image = torch.cat(reference_image, dim=0) + img_reference = reference_image[0] if isinstance(reference_image, list) else reference_image + reference_image_height, reference_image_width = self.image_processor.get_default_height_width(img_reference) + + reference_image_width = reference_image_width // multiple_of * multiple_of + reference_image_height = reference_image_height // multiple_of * multiple_of + reference_image = self.image_processor.resize( + reference_image, reference_image_height, reference_image_width + ) + reference_image = self.image_processor.preprocess( + reference_image, + reference_image_height, + reference_image_width, + crops_coords=crops_coords, + resize_mode=resize_mode, + ) + else: + pass + + reference_image_latents = None + reference_image_latent_ids = None + if reference_image is not None: + reference_image_latents, reference_image_latent_ids = self.prepare_image_latents( + images=[reference_image], + batch_size=batch_size * num_images_per_prompt, + generator=generator, + device=device, + dtype=self.vae.dtype, + ) + + if reference_image_latent_ids is not None: + latent_ids = torch.cat([latent_ids, reference_image_latent_ids], dim=1) + elif image_latent_ids is not None: + latent_ids = torch.cat([latent_ids, image_latent_ids], dim=1) + + mask = None + if mask_image is not None: + mask_condition = self.mask_processor.preprocess( + mask_image, height=height, width=width, resize_mode=resize_mode, crops_coords=crops_coords + ) + if init_image is not None: + masked_image = init_image * (mask_condition < 0.5).to(init_image.dtype) + else: + masked_image = None + mask, _ = self.prepare_mask_latents( + mask_condition, + masked_image, + batch_size, + self.latent_channels, + num_images_per_prompt, + height, + width, + prompt_embeds.dtype, + device, + generator, + ) # 6. Prepare timesteps sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas if hasattr(self.scheduler.config, "use_flow_sigmas") and self.scheduler.config.use_flow_sigmas: @@ -908,6 +1070,13 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi sigmas=sigmas, mu=mu, ) + if reference_image is not None or mask_image is not None: + timesteps, num_inference_steps = self.get_timesteps(num_inference_steps, strength, device) + if num_inference_steps < 1: + raise ValueError( + f"After adjusting the num_inference_steps by strength parameter: {strength}, the number of pipeline" + f"steps is {num_inference_steps} which is < 1 and not appropriate for this pipeline." + ) self._num_timesteps = len(timesteps) # 7. Denoising loop @@ -924,7 +1093,10 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi latent_model_input = latents.to(self.transformer.dtype) latent_image_ids = latent_ids - if image_latents is not None: + if reference_image_latents is not None: + latent_model_input = torch.cat([latents, reference_image_latents], dim=1) + latent_image_ids = latent_ids + elif image_latents is not None: latent_model_input = torch.cat([latents, image_latents], dim=1).to(self.transformer.dtype) latent_image_ids = torch.cat([latent_ids, image_latent_ids], dim=1) @@ -953,7 +1125,9 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi negative_kwargs = None # For editing pipelines, we need to slice the output to remove condition latents - output_slice = latents.size(1) if image_latents is not None else None + output_slice = ( + latents.size(1) if (image_latents is not None or reference_image_latents is not None) else None + ) noise_pred = self.predict_noise_maybe_with_cfg( do_true_cfg=self.do_classifier_free_guidance, @@ -964,9 +1138,22 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi output_slice=output_slice, ) + latents_dtype = latents.dtype # Compute the previous noisy sample x_t -> x_t-1 with automatic CFG sync latents = self.scheduler_step_maybe_with_cfg(noise_pred, t, latents, self.do_classifier_free_guidance) + if mask is not None and image_latents is not None: + init_latents_proper = image_latents + if i < len(timesteps) - 1: + noise_timestep = timesteps[i + 1] + init_latents_proper = self.scheduler.scale_noise( + init_latents_proper, torch.tensor([noise_timestep], device=device), latents + ) + latents = (1 - mask) * init_latents_proper + mask * latents + + if latents.dtype != latents_dtype and torch.backends.mps.is_available(): + latents = latents.to(latents_dtype) + if callback_on_step_end is not None: callback_kwargs = {} for k in callback_on_step_end_tensor_inputs: @@ -978,7 +1165,7 @@ class Flux2KleinPipeline(nn.Module, CFGParallelMixin, SupportImageInput, Diffusi self._current_timestep = None - latents = self._unpack_latents_with_ids(latents, latent_ids) + latents = self._unpack_latents_with_ids(latents, original_latent_ids) latents_bn_mean = self.vae.bn.running_mean.view(1, -1, 1, 1).to(latents.device, latents.dtype) latents_bn_std = torch.sqrt(self.vae.bn.running_var.view(1, -1, 1, 1) + self.vae.config.batch_norm_eps).to( diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 11ba59e43..1e4575836 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1426,6 +1426,9 @@ async def edit_images( background: str | None = Form("auto"), output_compression: Annotated[int, Form(ge=0, le=100)] = 100, user: str | None = Form(None), # unused now + # vllm-omni extensions for image editing + mask_image: str | UploadFile | None = None, + reference_image: str | UploadFile | None = None, # vllm-omni extensions for diffusion control negative_prompt: str | None = Form(None), num_inference_steps: int | None = Form(None), @@ -1479,6 +1482,14 @@ async def edit_images( prompt["multi_modal_data"] = {} prompt["multi_modal_data"]["image"] = pil_images + if mask_image is not None: + loaded = await _load_input_images([mask_image]) + prompt["multi_modal_data"]["mask_image"] = loaded[0] + + if reference_image is not None: + loaded = await _load_input_images([reference_image]) + prompt["multi_modal_data"]["reference_image"] = loaded[0] + # 3 Build sample params gen_params = OmniDiffusionSamplingParams() # 3.0 Init with system default values diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 34ddbbd30..28d6ef277 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -2197,7 +2197,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): except Exception as e: # pragma: no cover - safeguard logger.warning("Failed to parse LoRA request: %s", e) - # Add reference image if provided + # Add reference image if provided (from messages content) if pil_images: if len(pil_images) == 1: gen_prompt["multi_modal_data"] = {} From 1e3bb365214e4bd591e4a185b828f756d21b8fe7 Mon Sep 17 00:00:00 2001 From: Nick Cao Date: Thu, 16 Apr 2026 21:24:47 -0400 Subject: [PATCH 33/97] [Refactor] Remove sox from dependencies (#2745) Signed-off-by: Nick Cao Co-authored-by: Claude --- docker/Dockerfile.ci | 2 +- docker/Dockerfile.cuda | 2 +- docker/Dockerfile.rocm | 2 +- .../installation/gpu/rocm.inc.md | 2 +- .../examples/offline_inference/qwen3_tts.md | 4 ++-- .../offline_inference/qwen3_tts/README.md | 4 ++-- requirements/common.txt | 1 - tests/utils/test_audio.py | 22 ++++++++++++++++-- .../qwen3_tts/tokenizer_25hz/vq/speech_vq.py | 9 ++------ vllm_omni/utils/audio.py | 23 +++++++++++++++++++ 10 files changed, 53 insertions(+), 18 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 2a98de1b8..9cbf89d0b 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -7,7 +7,7 @@ COPY . . # Install system dependencies RUN apt-get update && \ - apt-get install -y espeak-ng git sox libsox-fmt-all jq && \ + apt-get install -y espeak-ng git jq && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 6ed5b7d27..28e10f4fb 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -7,7 +7,7 @@ WORKDIR ${COMMON_WORKDIR} # Step 1: Setup - Install system dependencies RUN apt-get update && \ - apt-get install -y git sox libsox-fmt-all jq && \ + apt-get install -y git jq && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index b34478389..a54aa3b79 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -21,7 +21,7 @@ WORKDIR ${COMMON_WORKDIR} # Need to include ffmpeg because vllm rocm upstream docker image # does not include it. RUN apt-get update && \ - apt-get install -y espeak-ng ffmpeg git sox libsox-fmt-all jq && \ + apt-get install -y espeak-ng ffmpeg git jq && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index 1a683d174..5dfea8d2f 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -26,7 +26,7 @@ uv pip install vllm-omni # Optional if want to run Qwen3 TTS uv pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm -uv pip install onnxruntime-rocm sox +uv pip install onnxruntime-rocm ``` # --8<-- [end:pre-built-wheels] diff --git a/docs/user_guide/examples/offline_inference/qwen3_tts.md b/docs/user_guide/examples/offline_inference/qwen3_tts.md index 19fea4132..4ece5219d 100644 --- a/docs/user_guide/examples/offline_inference/qwen3_tts.md +++ b/docs/user_guide/examples/offline_inference/qwen3_tts.md @@ -18,11 +18,11 @@ Please refer to the [stage configuration documentation](https://docs.vllm.ai/pro ### ROCm Dependencies -You will need to install these two dependencies `onnxruntime-rocm` and `sox`. +You will need to install the dependency `onnxruntime-rocm`. ``` pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm -pip install onnxruntime-rocm sox +pip install onnxruntime-rocm ``` ## Quick Start diff --git a/examples/offline_inference/qwen3_tts/README.md b/examples/offline_inference/qwen3_tts/README.md index bf59dc9ba..c38a2b462 100644 --- a/examples/offline_inference/qwen3_tts/README.md +++ b/examples/offline_inference/qwen3_tts/README.md @@ -15,11 +15,11 @@ Please refer to the [stage configuration documentation](https://docs.vllm.ai/pro ### ROCm Dependencies -You will need to install these two dependencies `onnxruntime-rocm` and `sox`. +You will need to install the dependency `onnxruntime-rocm`. ``` pip uninstall onnxruntime # should be removed before we can install onnxruntime-rocm -pip install onnxruntime-rocm sox +pip install onnxruntime-rocm ``` ## Quick Start diff --git a/requirements/common.txt b/requirements/common.txt index 1fff58444..1f44d343c 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -10,7 +10,6 @@ tqdm>=4.66.0 torchsde>=0.2.6 openai-whisper>=20250625 imageio[ffmpeg]>=2.37.2 -sox>=1.5.0 x-transformers>=2.12.2 einops>=0.8.1 prettytable>=3.8.0 diff --git a/tests/utils/test_audio.py b/tests/utils/test_audio.py index cfbd2501b..0e483e646 100644 --- a/tests/utils/test_audio.py +++ b/tests/utils/test_audio.py @@ -1,12 +1,13 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -"""Unit tests for vllm_omni.utils.audio.mel_filter_bank.""" +"""Unit tests for vllm_omni.utils.audio.""" +import numpy as np import pytest import torch -from vllm_omni.utils.audio import mel_filter_bank +from vllm_omni.utils.audio import mel_filter_bank, peak_normalize # Parameter combinations used across the codebase. _PARAM_SETS = [ @@ -59,3 +60,20 @@ class TestMelFilterBank: last_nonzero_low = (fb_low.sum(dim=0) > 0).nonzero()[-1].item() last_nonzero_high = (fb_high.sum(dim=0) > 0).nonzero()[-1].item() assert last_nonzero_high > last_nonzero_low + + +class TestPeakNormalize: + def test_silence_unchanged(self): + """All-zero input should remain all-zero.""" + audio = np.zeros(1600, dtype=np.float32) + result = peak_normalize(audio, db_level=-6.0) + np.testing.assert_array_equal(result, audio) + + def test_peak_reaches_target(self): + """After normalization, peak amplitude should be at target dB.""" + rng = np.random.default_rng(7) + audio = rng.uniform(-0.4, 0.4, size=16000).astype(np.float32) + + result = peak_normalize(audio, db_level=-6.0) + peak_db = 20 * np.log10(np.abs(result).max()) + np.testing.assert_allclose(peak_db, -6.0, atol=1e-4) diff --git a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py index 9bb2f78c5..92cecbff1 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py +++ b/vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/speech_vq.py @@ -17,14 +17,13 @@ import operator from itertools import accumulate import onnxruntime -import sox import torch import torch.nn as nn import torch.nn.functional as F import torchaudio.compliance.kaldi as kaldi from torch import Tensor -from vllm_omni.utils.audio import mel_filter_bank +from vllm_omni.utils.audio import mel_filter_bank, peak_normalize from .core_vq import DistributedGroupResidualVectorQuantization from .whisper_encoder import Conv1d, ConvTranspose1d, WhisperEncoder @@ -153,9 +152,6 @@ class XVectorExtractor(nn.Module): audio_codec_with_xvector, sess_options=option, providers=providers ) - self.tfm = sox.Transformer() - self.tfm.norm(db_level=-6) - self.mel_ext = MelSpectrogramFeatures( filter_length=1024, hop_length=160, @@ -183,8 +179,7 @@ class XVectorExtractor(nn.Module): return norm_embedding.numpy(), ref_mel.permute(0, 2, 1).squeeze(0).numpy() def sox_norm(self, audio): - wav_norm = self.tfm.build_array(input_array=audio, sample_rate_in=16000) - return wav_norm + return peak_normalize(audio, db_level=-6) class WhisperEncoderVQ(WhisperEncoder): diff --git a/vllm_omni/utils/audio.py b/vllm_omni/utils/audio.py index 490737bd5..cc25c1794 100644 --- a/vllm_omni/utils/audio.py +++ b/vllm_omni/utils/audio.py @@ -3,6 +3,7 @@ """Audio utility functions shared across models and entrypoints.""" +import numpy as np import torch from torchaudio.functional import melscale_fbanks @@ -43,3 +44,25 @@ def mel_filter_bank( mel_scale="slaney", norm="slaney", ).T + + +def peak_normalize( + audio: np.ndarray, + db_level: float = -6.0, +) -> np.ndarray: + """Normalize audio so peak amplitude reaches a target dB level. + + Drop-in replacement for ``sox.Transformer().norm(db_level=...)``. + + Args: + audio: Input waveform as a 1-D numpy array. + db_level: Target peak amplitude in dBFS. + + Returns: + Normalized waveform with the same dtype as *audio*. + """ + peak = np.abs(audio).max() + if peak == 0: + return audio + target = 10.0 ** (db_level / 20.0) + return audio * (target / peak) From 3079e943dea14380df088c6c1bbf108d97be824c Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Fri, 17 Apr 2026 10:40:36 +0800 Subject: [PATCH 34/97] [Bugfix] enforce max_sequence_length for Qwen-Image and Wan2.2 series before encoding (#2847) Signed-off-by: david6666666 <530634352@qq.com> --- .../test_qwen_image_max_sequence_length.py | 260 ++++++++++++++++++ .../wan2_2/test_wan22_max_sequence_length.py | 150 ++++++++++ .../models/qwen_image/pipeline_qwen_image.py | 44 ++- .../qwen_image/pipeline_qwen_image_edit.py | 49 +++- .../pipeline_qwen_image_edit_plus.py | 48 +++- .../qwen_image/pipeline_qwen_image_layered.py | 39 ++- .../models/wan2_2/pipeline_wan2_2.py | 46 +++- .../models/wan2_2/pipeline_wan2_2_i2v.py | 46 +++- .../models/wan2_2/pipeline_wan2_2_ti2v.py | 46 +++- .../models/wan2_2/pipeline_wan2_2_vace.py | 5 +- vllm_omni/diffusion/utils/prompt_utils.py | 38 +++ 11 files changed, 744 insertions(+), 27 deletions(-) create mode 100644 tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py create mode 100644 tests/diffusion/models/wan2_2/test_wan22_max_sequence_length.py create mode 100644 vllm_omni/diffusion/utils/prompt_utils.py diff --git a/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py new file mode 100644 index 000000000..f5676a005 --- /dev/null +++ b/tests/diffusion/models/qwen_image/test_qwen_image_max_sequence_length.py @@ -0,0 +1,260 @@ +import inspect +from types import SimpleNamespace + +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image import ( + QwenImagePipeline, +) +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit import ( + QwenImageEditPipeline, +) +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_edit_plus import ( + QwenImageEditPlusPipeline, +) +from vllm_omni.diffusion.models.qwen_image.pipeline_qwen_image_layered import ( + QwenImageLayeredPipeline, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _RejectingTextEncoder: + dtype = torch.float32 + + def __call__(self, *args, **kwargs): + raise AssertionError("text encoder should not run for prompts that exceed max_sequence_length") + + +class _FakeModelInputs: + def __init__(self, total_sequence_length: int): + attention_mask = torch.ones((1, total_sequence_length), dtype=torch.long) + self.input_ids = attention_mask.clone() + self.attention_mask = attention_mask + self.pixel_values = None + self.image_grid_thw = None + + def to(self, device): + return self + + +class _FakeTokenizer: + def __init__(self, total_sequence_length: int | list[int]): + if isinstance(total_sequence_length, list): + self.total_sequence_lengths = list(total_sequence_length) + else: + self.total_sequence_lengths = [total_sequence_length] + + def __call__(self, *args, **kwargs): + if len(self.total_sequence_lengths) > 1: + total_sequence_length = self.total_sequence_lengths.pop(0) + else: + total_sequence_length = self.total_sequence_lengths[0] + return _FakeModelInputs(total_sequence_length) + + +class _FakeProcessor(_FakeTokenizer): + pass + + +class _FakeScheduler: + def __init__(self): + self.begin_index = None + + def set_begin_index(self, begin_index: int): + self.begin_index = begin_index + + +PIPELINE_CASES = [ + pytest.param(QwenImagePipeline, 34, "tokenizer", id="qwen-image"), + pytest.param(QwenImageLayeredPipeline, 34, "tokenizer", id="qwen-image-layered"), + pytest.param(QwenImageEditPipeline, 64, "processor", id="qwen-image-edit"), + pytest.param(QwenImageEditPlusPipeline, 64, "processor", id="qwen-image-edit-plus"), +] + + +def _make_pipeline( + pipeline_class: type, + *, + total_sequence_length: int, + drop_idx: int, + input_kind: str, +): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.text_encoder = _RejectingTextEncoder() + pipeline.tokenizer_max_length = 1024 + pipeline.prompt_template_encode = "{}" + pipeline.prompt_template_encode_start_idx = drop_idx + pipeline.tokenizer = _FakeTokenizer([total_sequence_length, 0]) + if input_kind == "processor": + pipeline.processor = _FakeProcessor(total_sequence_length) + return pipeline + + +@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES) +def test_encode_prompt_rejects_prompt_longer_than_default_max_sequence_length( + pipeline_class: type, + drop_idx: int, + input_kind: str, +): + pipeline = _make_pipeline( + pipeline_class, + total_sequence_length=1025, + drop_idx=drop_idx, + input_kind=input_kind, + ) + + with pytest.raises(ValueError, match=r"got 1025 tokens, but `max_sequence_length` is 1024"): + pipeline.encode_prompt(prompt="prompt") + + +@pytest.mark.parametrize(("pipeline_class", "drop_idx", "input_kind"), PIPELINE_CASES) +def test_encode_prompt_rejects_prompt_longer_than_explicit_max_sequence_length( + pipeline_class: type, + drop_idx: int, + input_kind: str, +): + pipeline = _make_pipeline( + pipeline_class, + total_sequence_length=17, + drop_idx=drop_idx, + input_kind=input_kind, + ) + + with pytest.raises(ValueError, match=r"got 17 tokens, but `max_sequence_length` is 16"): + pipeline.encode_prompt(prompt="prompt", max_sequence_length=16) + + +def test_prepare_encode_defaults_to_tokenizer_max_length(): + pipeline = object.__new__(QwenImagePipeline) + nn.Module.__init__(pipeline) + pipeline.tokenizer_max_length = 1024 + pipeline.vae_scale_factor = 8 + pipeline.default_sample_size = 128 + pipeline.scheduler = _FakeScheduler() + pipeline._extract_prompts = lambda prompts: (["prompt"], None) + + captured = {} + + def _fake_prepare_generation_context(**kwargs): + captured["max_sequence_length"] = kwargs["max_sequence_length"] + embeds = torch.ones((1, 1, 1)) + mask = torch.ones((1, 1), dtype=torch.long) + return { + "prompt_embeds": embeds, + "prompt_embeds_mask": mask, + "negative_prompt_embeds": None, + "negative_prompt_embeds_mask": None, + "latents": embeds, + "timesteps": torch.tensor([1]), + "do_true_cfg": False, + "guidance": None, + "img_shapes": [[(1, 1, 1)]], + "txt_seq_lens": [1], + "negative_txt_seq_lens": None, + } + + pipeline._prepare_generation_context = _fake_prepare_generation_context + state = SimpleNamespace( + prompts=["prompt"], + sampling=SimpleNamespace( + height=None, + width=None, + num_inference_steps=None, + sigmas=None, + guidance_scale_provided=False, + num_outputs_per_prompt=0, + generator=None, + true_cfg_scale=None, + max_sequence_length=None, + ), + ) + + pipeline.prepare_encode(state) + + assert captured["max_sequence_length"] == 1024 + + +@pytest.mark.parametrize( + ("pipeline_class", "drop_idx"), + [ + pytest.param(QwenImageEditPipeline, 64, id="qwen-image-edit"), + pytest.param(QwenImageEditPlusPipeline, 64, id="qwen-image-edit-plus"), + ], +) +def test_edit_pipelines_validate_text_prompt_length_before_image_token_expansion( + pipeline_class: type, + drop_idx: int, +): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.text_encoder = _RejectingTextEncoder() + pipeline.tokenizer_max_length = 1024 + pipeline.prompt_template_encode = "{}" + pipeline.prompt_template_encode_start_idx = drop_idx + pipeline.tokenizer = _FakeTokenizer([8, 0]) + pipeline.processor = _FakeProcessor(drop_idx + 1500) + + with pytest.raises(AssertionError, match="text encoder should not run"): + pipeline.encode_prompt(prompt="short prompt") + + +@pytest.mark.parametrize( + "pipeline_class", + [ + pytest.param(QwenImagePipeline, id="qwen-image"), + pytest.param(QwenImageLayeredPipeline, id="qwen-image-layered"), + ], +) +def test_qwen_generation_validator_excludes_template_suffix_from_budget(pipeline_class: type): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.text_encoder = _RejectingTextEncoder() + pipeline.tokenizer_max_length = 1024 + pipeline.prompt_template_encode = "{}" + pipeline.prompt_template_encode_start_idx = 34 + pipeline.tokenizer = _FakeTokenizer([1029, 5]) + + with pytest.raises(AssertionError, match="text encoder should not run"): + pipeline.encode_prompt(prompt="boundary prompt") + + +@pytest.mark.parametrize( + "pipeline_class", + [ + pytest.param(QwenImageEditPipeline, id="qwen-image-edit"), + pytest.param(QwenImageEditPlusPipeline, id="qwen-image-edit-plus"), + ], +) +def test_qwen_edit_validator_excludes_image_placeholders_from_budget(pipeline_class: type): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.text_encoder = _RejectingTextEncoder() + pipeline.tokenizer_max_length = 1024 + pipeline.prompt_template_encode = "{}" + pipeline.prompt_template_encode_start_idx = 64 + pipeline.tokenizer = _FakeTokenizer([30, 20]) + pipeline.processor = _FakeProcessor(1500) + + with pytest.raises(AssertionError, match="text encoder should not run"): + pipeline.encode_prompt(prompt="short prompt") + + +@pytest.mark.parametrize( + "pipeline_class", + [ + QwenImagePipeline, + QwenImageLayeredPipeline, + QwenImageEditPipeline, + QwenImageEditPlusPipeline, + ], +) +def test_forward_max_sequence_length_default_is_1024(pipeline_class: type): + assert inspect.signature(pipeline_class.forward).parameters["max_sequence_length"].default == 1024 diff --git a/tests/diffusion/models/wan2_2/test_wan22_max_sequence_length.py b/tests/diffusion/models/wan2_2/test_wan22_max_sequence_length.py new file mode 100644 index 000000000..64c2b271c --- /dev/null +++ b/tests/diffusion/models/wan2_2/test_wan22_max_sequence_length.py @@ -0,0 +1,150 @@ +from types import SimpleNamespace + +import PIL.Image +import pytest +import torch +from torch import nn + +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + WAN22_MAX_SEQUENCE_LENGTH, + Wan22Pipeline, +) +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_i2v import ( + Wan22I2VPipeline, +) +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_ti2v import ( + Wan22TI2VPipeline, +) +from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2_vace import ( + Wan22VACEPipeline, +) + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +class _RejectingTextEncoder: + dtype = torch.float32 + + def __call__(self, *args, **kwargs): + raise AssertionError("text encoder should not run for prompts that exceed max_sequence_length") + + +class _FakeTokenBatch: + def __init__(self, total_sequence_length: int): + attention_mask = torch.ones((1, total_sequence_length), dtype=torch.long) + self.input_ids = attention_mask.clone() + self.attention_mask = attention_mask + + +class _FakeTokenizer: + def __init__(self, total_sequence_length: int): + self.total_sequence_length = total_sequence_length + + def __call__(self, *args, **kwargs): + return _FakeTokenBatch(self.total_sequence_length) + + +PIPELINE_CASES = [ + pytest.param(Wan22Pipeline, id="wan22-t2v"), + pytest.param(Wan22I2VPipeline, id="wan22-i2v"), + pytest.param(Wan22TI2VPipeline, id="wan22-ti2v"), + pytest.param(Wan22VACEPipeline, id="wan22-vace"), +] + + +def _make_pipeline(pipeline_class: type, *, total_sequence_length: int): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.device = torch.device("cpu") + pipeline.text_encoder = _RejectingTextEncoder() + pipeline.tokenizer = _FakeTokenizer(total_sequence_length) + pipeline.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH + return pipeline + + +@pytest.mark.parametrize("pipeline_class", PIPELINE_CASES) +def test_encode_prompt_rejects_prompt_longer_than_default_max_sequence_length(pipeline_class: type): + pipeline = _make_pipeline(pipeline_class, total_sequence_length=WAN22_MAX_SEQUENCE_LENGTH + 1) + + with pytest.raises(ValueError, match=r"got 513 tokens, but `max_sequence_length` is 512"): + pipeline.encode_prompt(prompt="prompt") + + +@pytest.mark.parametrize("pipeline_class", PIPELINE_CASES) +def test_encode_prompt_rejects_prompt_longer_than_explicit_max_sequence_length(pipeline_class: type): + pipeline = _make_pipeline(pipeline_class, total_sequence_length=17) + + with pytest.raises(ValueError, match=r"got 17 tokens, but `max_sequence_length` is 16"): + pipeline.encode_prompt(prompt="prompt", max_sequence_length=16) + + +def _sampling_params(**overrides): + defaults = dict( + height=None, + width=None, + num_frames=None, + num_inference_steps=None, + generator=None, + guidance_scale_provided=False, + guidance_scale_2=None, + boundary_ratio=None, + num_outputs_per_prompt=0, + max_sequence_length=None, + seed=None, + extra_args={}, + prompt_embeds=None, + negative_prompt_embeds=None, + ) + defaults.update(overrides) + return SimpleNamespace(**defaults) + + +@pytest.mark.parametrize( + ("pipeline_class", "prompt_value", "forward_kwargs"), + [ + pytest.param(Wan22Pipeline, "prompt", {}, id="wan22-t2v"), + pytest.param( + Wan22I2VPipeline, + {"prompt": "prompt", "multi_modal_data": {"image": PIL.Image.new("RGB", (64, 64))}}, + {"image": PIL.Image.new("RGB", (64, 64))}, + id="wan22-i2v", + ), + pytest.param( + Wan22TI2VPipeline, + {"prompt": "prompt", "multi_modal_data": {"image": PIL.Image.new("RGB", (64, 64))}}, + {"image": PIL.Image.new("RGB", (64, 64))}, + id="wan22-ti2v", + ), + pytest.param(Wan22VACEPipeline, "prompt", {}, id="wan22-vace"), + ], +) +def test_forward_defaults_to_wan22_tokenizer_max_length( + pipeline_class: type, + prompt_value, + forward_kwargs, +): + pipeline = object.__new__(pipeline_class) + nn.Module.__init__(pipeline) + pipeline.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH + pipeline.boundary_ratio = None + pipeline.vae_scale_factor_temporal = 4 + pipeline.vae_scale_factor_spatial = 8 + pipeline.transformer_config = SimpleNamespace(patch_size=(1, 2, 2)) + + captured = {} + + def _fake_check_inputs(*args, **kwargs): + captured["max_sequence_length"] = kwargs["max_sequence_length"] + raise RuntimeError("stop after capture") + + pipeline.check_inputs = _fake_check_inputs + + req = SimpleNamespace( + prompts=[prompt_value], + sampling_params=_sampling_params(), + ) + + with pytest.raises(RuntimeError, match="stop after capture"): + pipeline.forward(req, **forward_kwargs) + + assert captured["max_sequence_length"] == WAN22_MAX_SEQUENCE_LENGTH diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py index 9f75c8453..9ef0cacd5 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image.py @@ -34,6 +34,9 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.prompt_utils import ( + validate_prompt_sequence_lengths, +) from vllm_omni.diffusion.utils.size_utils import ( normalize_min_aligned_size, ) @@ -363,8 +366,10 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP "that was used to generate `negative_prompt_embeds`." ) - if max_sequence_length is not None and max_sequence_length > 1024: - raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -378,6 +383,8 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP self, prompt: str | list[str] = None, dtype: torch.dtype | None = None, + max_sequence_length: int | None = None, + prompt_name: str = "prompt", ): dtype = dtype or self.text_encoder.dtype @@ -388,12 +395,27 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP txt = [template.format(e) for e in prompt] txt_tokens = self.tokenizer( txt, - max_length=self.tokenizer_max_length + drop_idx, padding=True, - truncation=True, + truncation=False, return_tensors="pt", ).to(self.device) - # print(f"attention mask: {txt_tokens.attention_mask}") + # Validate only the user prompt contribution. The Qwen template also + # adds a fixed suffix after the user text, so subtracting only + # prompt_template_encode_start_idx would overcount near-limit prompts. + template_tokens = self.tokenizer( + [template.format("")], + padding=True, + truncation=False, + return_tensors="pt", + ).to(self.device) + validate_prompt_sequence_lengths( + txt_tokens.attention_mask, + max_sequence_length=max_sequence_length or self.tokenizer_max_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name=prompt_name, + baseline_attention_mask=template_tokens.attention_mask, + error_context="after applying the Qwen prompt template", + ) encoder_hidden_states = self.text_encoder( input_ids=txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask, @@ -422,6 +444,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, + prompt_name: str = "prompt", ): r""" @@ -439,7 +462,11 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] if prompt_embeds is None: - prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt) + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + prompt, + max_sequence_length=max_sequence_length, + prompt_name=prompt_name, + ) prompt_embeds = prompt_embeds[:, :max_sequence_length] prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] @@ -632,6 +659,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP prompt_embeds_mask=negative_prompt_embeds_mask, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, + prompt_name="negative_prompt", ) else: negative_prompt_embeds = None @@ -703,7 +731,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP num_images_per_prompt=sampling.num_outputs_per_prompt if sampling.num_outputs_per_prompt > 0 else 1, generator=sampling.generator, true_cfg_scale=sampling.true_cfg_scale or 4.0, - max_sequence_length=sampling.max_sequence_length or 512, + max_sequence_length=sampling.max_sequence_length or self.tokenizer_max_length, attention_kwargs=kwargs.get("attention_kwargs"), ) @@ -934,7 +962,7 @@ class QwenImagePipeline(nn.Module, QwenImageCFGParallelMixin, DiffusionPipelineP output_type: str | None = "pil", attention_kwargs: dict[str, Any] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, + max_sequence_length: int = 1024, ) -> DiffusionOutput: extracted_prompt, negative_prompt = self._extract_prompts(req.prompts) prompt = extracted_prompt or prompt diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py index dd77d71b1..cef7fe473 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit.py @@ -37,6 +37,9 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.prompt_utils import ( + validate_prompt_sequence_lengths, +) from vllm_omni.diffusion.utils.size_utils import ( normalize_min_aligned_size, ) @@ -323,8 +326,10 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi "that was used to generate `negative_prompt_embeds`." ) - if max_sequence_length is not None and max_sequence_length > 1024: - raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -384,6 +389,8 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi prompt: str | list[str] = None, image: PIL.Image.Image | torch.Tensor | None = None, dtype: torch.dtype | None = None, + max_sequence_length: int | None = None, + prompt_name: str = "prompt", ): """Get prompt embeddings with image support for editing.""" dtype = dtype or self.text_encoder.dtype @@ -393,6 +400,33 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi template = self.prompt_template_encode drop_idx = self.prompt_template_encode_start_idx txt = [template.format(e) for e in prompt] + txt_tokens = self.tokenizer( + txt, + padding=True, + truncation=False, + return_tensors="pt", + ).to(self.device) + # The edit template contains fixed multimodal scaffolding around the + # instruction. Validate against the empty-template baseline so image + # placeholder text does not consume the user's text budget. + template_tokens = self.tokenizer( + [template.format("")], + padding=True, + truncation=False, + return_tensors="pt", + ).to(self.device) + # Qwen-Image-Edit expands image placeholders into many vision tokens + # inside the processor. `max_sequence_length` is meant to constrain the + # prompt text length, so validate on the text template before image + # token expansion. + validate_prompt_sequence_lengths( + txt_tokens.attention_mask, + max_sequence_length=max_sequence_length or self.tokenizer_max_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name=prompt_name, + baseline_attention_mask=template_tokens.attention_mask, + error_context="after applying the Qwen prompt template", + ) # Use processor to handle both text and image inputs model_inputs = self.processor( @@ -434,6 +468,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, + prompt_name: str = "prompt", ): r""" @@ -453,7 +488,12 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] if prompt_embeds is None: - prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image) + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + prompt, + image, + max_sequence_length=max_sequence_length, + prompt_name=prompt_name, + ) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) @@ -624,7 +664,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi output_type: str | None = "pil", attention_kwargs: dict[str, Any] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, + max_sequence_length: int = 1024, ) -> DiffusionOutput: """Forward pass for image editing.""" # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") @@ -739,6 +779,7 @@ class QwenImageEditPipeline(nn.Module, SupportImageInput, QwenImageCFGParallelMi prompt_embeds_mask=negative_prompt_embeds_mask, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, + prompt_name="negative_prompt", ) num_channels_latents = self.transformer.in_channels // 4 diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py index 6f6c9d2ba..a2702f3d2 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_edit_plus.py @@ -40,6 +40,9 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.prompt_utils import ( + validate_prompt_sequence_lengths, +) from vllm_omni.diffusion.utils.size_utils import ( normalize_min_aligned_size, ) @@ -283,8 +286,10 @@ class QwenImageEditPlusPipeline( "that was used to generate `negative_prompt_embeds`." ) - if max_sequence_length is not None and max_sequence_length > 1024: - raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -299,6 +304,8 @@ class QwenImageEditPlusPipeline( prompt: str | list[str], image: list[torch.Tensor] | torch.Tensor | None = None, dtype: torch.dtype | None = None, + max_sequence_length: int | None = None, + prompt_name: str = "prompt", ): """Get prompt embeddings with support for multiple images.""" dtype = dtype or self.text_encoder.dtype @@ -319,6 +326,32 @@ class QwenImageEditPlusPipeline( template = self.prompt_template_encode drop_idx = self.prompt_template_encode_start_idx txt = [template.format(base_img_prompt + e) for e in prompt] + txt_tokens = self.tokenizer( + txt, + padding=True, + truncation=False, + return_tensors="pt", + ).to(self.device) + # Multi-image edit prepends "Picture N" placeholders before the user + # instruction. Subtract the placeholder-aware baseline so attached + # images do not reduce the remaining prompt budget. + template_tokens = self.tokenizer( + [template.format(base_img_prompt)], + padding=True, + truncation=False, + return_tensors="pt", + ).to(self.device) + # The processor expands image placeholders into many vision tokens. + # `max_sequence_length` should guard the prompt text length before that + # multimodal expansion happens. + validate_prompt_sequence_lengths( + txt_tokens.attention_mask, + max_sequence_length=max_sequence_length or self.tokenizer_max_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name=prompt_name, + baseline_attention_mask=template_tokens.attention_mask, + error_context="after applying the Qwen prompt template", + ) # Use processor to handle both text and image inputs model_inputs = self.processor( @@ -360,6 +393,7 @@ class QwenImageEditPlusPipeline( prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, + prompt_name: str = "prompt", ): r""" @@ -379,7 +413,12 @@ class QwenImageEditPlusPipeline( batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] if prompt_embeds is None: - prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt, image) + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + prompt, + image, + max_sequence_length=max_sequence_length, + prompt_name=prompt_name, + ) _, seq_len, _ = prompt_embeds.shape prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) @@ -557,7 +596,7 @@ class QwenImageEditPlusPipeline( output_type: str | None = "pil", attention_kwargs: dict[str, Any] | None = None, callback_on_step_end_tensor_inputs: list[str] = ["latents"], - max_sequence_length: int = 512, + max_sequence_length: int = 1024, ) -> DiffusionOutput: """Forward pass for image editing with support for multiple images.""" # TODO: In online mode, sometimes it receives [{"negative_prompt": None}, {...}], so cannot use .get("...", "") @@ -692,6 +731,7 @@ class QwenImageEditPlusPipeline( prompt_embeds_mask=negative_prompt_embeds_mask, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, + prompt_name="negative_prompt", ) num_channels_latents = self.transformer.in_channels // 4 diff --git a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py index 38866d89c..905ef5b42 100644 --- a/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py +++ b/vllm_omni/diffusion/models/qwen_image/pipeline_qwen_image_layered.py @@ -36,6 +36,9 @@ from vllm_omni.diffusion.models.qwen_image.qwen_image_transformer import ( ) from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.prompt_utils import ( + validate_prompt_sequence_lengths, +) from vllm_omni.diffusion.utils.size_utils import ( normalize_min_aligned_size, ) @@ -340,8 +343,10 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as "generate `negative_prompt_embeds`." ) - if max_sequence_length is not None and max_sequence_length > 1024: - raise ValueError(f"`max_sequence_length` cannot be greater than 1024 but is {max_sequence_length}") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) def _extract_masked_hidden(self, hidden_states: torch.Tensor, mask: torch.Tensor): bool_mask = mask.bool() @@ -356,6 +361,8 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as prompt: str | list[str] | None = None, device: torch.device | None = None, dtype: torch.dtype | None = None, + max_sequence_length: int | None = None, + prompt_name: str = "prompt", ): device = device or self.device dtype = dtype or self.text_encoder.dtype @@ -368,8 +375,26 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as txt_tokens = self.tokenizer( txt, padding=True, + truncation=False, return_tensors="pt", ).to(device) + # The layered template also appends fixed non-user tokens after the + # editable text, so use the empty-template tokenized baseline instead of + # counting everything after prompt_template_encode_start_idx. + template_tokens = self.tokenizer( + [template.format("")], + padding=True, + truncation=False, + return_tensors="pt", + ).to(device) + validate_prompt_sequence_lengths( + txt_tokens.attention_mask, + max_sequence_length=max_sequence_length or self.tokenizer_max_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name=prompt_name, + baseline_attention_mask=template_tokens.attention_mask, + error_context="after applying the Qwen prompt template", + ) encoder_hidden_states = self.text_encoder( input_ids=txt_tokens.input_ids, attention_mask=txt_tokens.attention_mask, @@ -399,6 +424,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as prompt_embeds: torch.Tensor | None = None, prompt_embeds_mask: torch.Tensor | None = None, max_sequence_length: int = 1024, + prompt_name: str = "prompt", ): r""" @@ -419,7 +445,11 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] if prompt_embeds is None: - prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds(prompt) + prompt_embeds, prompt_embeds_mask = self._get_qwen_prompt_embeds( + prompt, + max_sequence_length=max_sequence_length, + prompt_name=prompt_name, + ) prompt_embeds = prompt_embeds[:, :max_sequence_length] prompt_embeds_mask = prompt_embeds_mask[:, :max_sequence_length] @@ -603,7 +633,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as negative_prompt_embeds_mask: torch.Tensor | None = None, output_type: str | None = "pil", attention_kwargs: dict[str, Any] | None = None, - max_sequence_length: int = 512, + max_sequence_length: int = 1024, resolution: int = 640, cfg_normalize: bool = False, use_en_prompt: bool = False, @@ -736,6 +766,7 @@ the image\n<|vision_start|><|image_pad|><|vision_end|><|im_end|>\n<|im_start|>as device=self.device, num_images_per_prompt=num_images_per_prompt, max_sequence_length=max_sequence_length, + prompt_name="negative_prompt", ) # 4. Prepare latent variables diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py index a1b10439c..e1249e889 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2.py @@ -29,12 +29,16 @@ from vllm_omni.diffusion.models.wan2_2.wan2_2_transformer import WanTransformer3 from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.prompt_utils import ( + validate_prompt_sequence_lengths, +) from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform logger = logging.getLogger(__name__) DEBUG_PERF = False WAN_SAMPLE_SOLVER_CHOICES = {"unipc", "euler"} +WAN22_MAX_SEQUENCE_LENGTH = 512 def build_wan_scheduler(sample_solver: str, flow_shift: float) -> Any: @@ -289,6 +293,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe pass self.boundary_ratio = od_config.boundary_ratio + self.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH # Determine which transformers to load based on boundary_ratio # boundary_ratio=1.0: only load transformer_2 (low-noise stage only) @@ -464,6 +469,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe negative_prompt_embeds=negative_prompt_embeds, guidance_scale_2=guidance_high if boundary_ratio is not None else None, boundary_ratio=boundary_ratio, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, ) if num_frames % self.vae_scale_factor_temporal != 1: @@ -497,7 +503,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe negative_prompt=negative_prompt, do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, device=device, dtype=dtype, ) @@ -801,6 +807,20 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe prompt = [prompt] if isinstance(prompt, str) else prompt prompt_clean = [self._prompt_clean(p) for p in prompt] batch_size = len(prompt_clean) + text_inputs_untruncated = self.tokenizer( + prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + error_context="for Wan2.2 text encoding", + ) text_inputs = self.tokenizer( prompt_clean, @@ -829,8 +849,24 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe if do_classifier_free_guidance: negative_prompt = negative_prompt or "" negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_clean = [self._prompt_clean(p) for p in negative_prompt] + neg_text_inputs_untruncated = self.tokenizer( + negative_prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + neg_text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name="negative_prompt", + error_context="for Wan2.2 text encoding", + ) neg_text_inputs = self.tokenizer( - [self._prompt_clean(p) for p in negative_prompt], + negative_prompt_clean, padding="max_length", max_length=max_sequence_length, truncation=True, @@ -902,6 +938,7 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe negative_prompt_embeds=None, guidance_scale_2=None, boundary_ratio=None, + max_sequence_length=None, ): if height % 16 != 0 or width % 16 != 0: raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") @@ -928,5 +965,10 @@ class Wan22Pipeline(nn.Module, CFGParallelMixin, ProgressBarMixin, DiffusionPipe ): raise ValueError(f"`negative_prompt` has to be of type `str` or `list` but is {type(negative_prompt)}") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) + if boundary_ratio is None and guidance_scale_2 is not None: raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.") diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py index ddc6e0bc2..ca042ca22 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_i2v.py @@ -25,6 +25,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin, _is_rank_zero from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + WAN22_MAX_SEQUENCE_LENGTH, build_wan_scheduler, create_transformer_from_config, load_transformer_config, @@ -35,6 +36,9 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.prompt_utils import ( + validate_prompt_sequence_lengths, +) from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -213,6 +217,7 @@ class Wan22I2VPipeline( # Text encoder self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH self.text_encoder = UMT5EncoderModel.from_pretrained( model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) @@ -392,6 +397,7 @@ class Wan22I2VPipeline( image_embeds=image_embeds, guidance_scale_2=guidance_high if boundary_ratio is not None else None, boundary_ratio=boundary_ratio, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, ) # Adjust num_frames to be compatible with VAE temporal scaling @@ -420,7 +426,7 @@ class Wan22I2VPipeline( negative_prompt=negative_prompt, do_classifier_free_guidance=guidance_low > 1.0 or guidance_high > 1.0, num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, device=device, dtype=dtype, ) @@ -671,6 +677,20 @@ class Wan22I2VPipeline( prompt = [prompt] if isinstance(prompt, str) else prompt prompt_clean = [self._prompt_clean(p) for p in prompt] batch_size = len(prompt_clean) + text_inputs_untruncated = self.tokenizer( + prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + error_context="for Wan2.2 text encoding", + ) text_inputs = self.tokenizer( prompt_clean, @@ -699,8 +719,24 @@ class Wan22I2VPipeline( if do_classifier_free_guidance: negative_prompt = negative_prompt or "" negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_clean = [self._prompt_clean(p) for p in negative_prompt] + neg_text_inputs_untruncated = self.tokenizer( + negative_prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + neg_text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name="negative_prompt", + error_context="for Wan2.2 text encoding", + ) neg_text_inputs = self.tokenizer( - [self._prompt_clean(p) for p in negative_prompt], + negative_prompt_clean, padding="max_length", max_length=max_sequence_length, truncation=True, @@ -842,6 +878,7 @@ class Wan22I2VPipeline( image_embeds=None, guidance_scale_2=None, boundary_ratio=None, + max_sequence_length=None, ): if image is None and image_embeds is None: raise ValueError("Provide either `image` or `image_embeds`. Cannot leave both undefined.") @@ -863,6 +900,11 @@ class Wan22I2VPipeline( if prompt is None and prompt_embeds is None: raise ValueError("Provide either `prompt` or `prompt_embeds`.") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) + if boundary_ratio is None and guidance_scale_2 is not None: raise ValueError("`guidance_scale_2` is only supported when `boundary_ratio` is set.") diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py index 62df13cbd..6cbd6d2d6 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_ti2v.py @@ -37,6 +37,7 @@ from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineL from vllm_omni.diffusion.models.interface import SupportImageInput from vllm_omni.diffusion.models.progress_bar import ProgressBarMixin from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( + WAN22_MAX_SEQUENCE_LENGTH, build_wan_scheduler, create_transformer_from_config, load_transformer_config, @@ -46,6 +47,9 @@ from vllm_omni.diffusion.models.wan2_2.pipeline_wan2_2 import ( ) from vllm_omni.diffusion.postprocess import interpolate_video_tensor from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.prompt_utils import ( + validate_prompt_sequence_lengths, +) from vllm_omni.inputs.data import OmniTextPrompt from vllm_omni.platforms import current_omni_platform @@ -185,6 +189,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress # Text encoder self.tokenizer = AutoTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) + self.tokenizer_max_length = WAN22_MAX_SEQUENCE_LENGTH self.text_encoder = UMT5EncoderModel.from_pretrained( model, subfolder="text_encoder", torch_dtype=dtype, local_files_only=local_files_only ).to(self.device) @@ -301,6 +306,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress width=width, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, ) # Adjust num_frames to be compatible with VAE temporal scaling @@ -324,7 +330,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress negative_prompt=negative_prompt, do_classifier_free_guidance=guidance_scale > 1.0, num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, device=device, dtype=dtype, ) @@ -518,6 +524,20 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress prompt = [prompt] if isinstance(prompt, str) else prompt prompt_clean = [self._prompt_clean(p) for p in prompt] batch_size = len(prompt_clean) + text_inputs_untruncated = self.tokenizer( + prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + error_context="for Wan2.2 text encoding", + ) text_inputs = self.tokenizer( prompt_clean, @@ -546,8 +566,24 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress if do_classifier_free_guidance: negative_prompt = negative_prompt or "" negative_prompt = batch_size * [negative_prompt] if isinstance(negative_prompt, str) else negative_prompt + negative_prompt_clean = [self._prompt_clean(p) for p in negative_prompt] + neg_text_inputs_untruncated = self.tokenizer( + negative_prompt_clean, + padding=True, + truncation=False, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + validate_prompt_sequence_lengths( + neg_text_inputs_untruncated.attention_mask, + max_sequence_length=max_sequence_length, + supported_max_sequence_length=self.tokenizer_max_length, + prompt_name="negative_prompt", + error_context="for Wan2.2 text encoding", + ) neg_text_inputs = self.tokenizer( - [self._prompt_clean(p) for p in negative_prompt], + negative_prompt_clean, padding="max_length", max_length=max_sequence_length, truncation=True, @@ -672,6 +708,7 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress width, prompt_embeds=None, negative_prompt_embeds=None, + max_sequence_length=None, ): if height % 16 != 0 or width % 16 != 0: raise ValueError(f"`height` and `width` have to be divisible by 16 but are {height} and {width}.") @@ -687,6 +724,11 @@ class Wan22TI2VPipeline(nn.Module, SupportImageInput, CFGParallelMixin, Progress if prompt is None and prompt_embeds is None: raise ValueError("Provide either `prompt` or `prompt_embeds`.") + if max_sequence_length is not None and max_sequence_length > self.tokenizer_max_length: + raise ValueError( + f"`max_sequence_length` cannot be greater than {self.tokenizer_max_length} but is {max_sequence_length}" + ) + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights using AutoWeightsLoader for vLLM integration.""" loader = AutoWeightsLoader(self) diff --git a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py index ea5233631..b661108cc 100644 --- a/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py +++ b/vllm_omni/diffusion/models/wan2_2/pipeline_wan2_2_vace.py @@ -187,6 +187,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): video=None, mask=None, reference_images=None, + max_sequence_length=None, ): super().check_inputs( prompt=prompt, @@ -195,6 +196,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): width=width, prompt_embeds=prompt_embeds, negative_prompt_embeds=negative_prompt_embeds, + max_sequence_length=max_sequence_length, ) # VACE-specific: validate video/mask/reference_images consistency @@ -491,6 +493,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): video=source_video, mask=source_mask, reference_images=reference_images, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, ) device = self.device @@ -509,7 +512,7 @@ class Wan22VACEPipeline(Wan22Pipeline, SupportImageInput): negative_prompt=negative_prompt, do_classifier_free_guidance=guidance_scale > 1.0, num_videos_per_prompt=req.sampling_params.num_outputs_per_prompt or 1, - max_sequence_length=req.sampling_params.max_sequence_length or 512, + max_sequence_length=req.sampling_params.max_sequence_length or self.tokenizer_max_length, device=device, dtype=dtype, ) diff --git a/vllm_omni/diffusion/utils/prompt_utils.py b/vllm_omni/diffusion/utils/prompt_utils.py new file mode 100644 index 000000000..fc1769f4d --- /dev/null +++ b/vllm_omni/diffusion/utils/prompt_utils.py @@ -0,0 +1,38 @@ +import torch + + +def validate_prompt_sequence_lengths( + attention_mask: torch.Tensor, + *, + max_sequence_length: int, + supported_max_sequence_length: int, + prompt_name: str = "prompt", + length_offset: int = 0, + baseline_attention_mask: torch.Tensor | None = None, + error_context: str, +) -> None: + sequence_lengths = attention_mask.sum(dim=1) + if baseline_attention_mask is not None: + # Some callers need to validate only the user-controlled portion of a + # templated prompt. In those cases we subtract the fully-tokenized + # template baseline instead of only removing a fixed prefix length, + # because the template may also contribute a suffix or image markers. + baseline_lengths = baseline_attention_mask.sum(dim=1) + if baseline_lengths.shape[0] == 1 and sequence_lengths.shape[0] > 1: + baseline_lengths = baseline_lengths.expand(sequence_lengths.shape[0]) + sequence_lengths = sequence_lengths - baseline_lengths + if length_offset: + sequence_lengths = sequence_lengths - length_offset + sequence_lengths = torch.clamp(sequence_lengths, min=0) + too_long = torch.nonzero(sequence_lengths > max_sequence_length, as_tuple=False) + if too_long.numel() == 0: + return + + batch_idx = int(too_long[0].item()) + actual_length = int(sequence_lengths[batch_idx].item()) + prompt_ref = f"`{prompt_name}` at batch index {batch_idx}" if attention_mask.shape[0] > 1 else f"`{prompt_name}`" + raise ValueError( + f"{prompt_ref} is too long {error_context}: got {actual_length} tokens, but " + f"`max_sequence_length` is {max_sequence_length}. Shorten the prompt or increase " + f"`max_sequence_length` up to {supported_max_sequence_length}." + ) From 12378824fbc3ecb8c9c33718ff1a861786331e9a Mon Sep 17 00:00:00 2001 From: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Date: Fri, 17 Apr 2026 12:07:10 +0800 Subject: [PATCH 35/97] [Bugfix] Preserve default diffusion sampling params in default stage (#2780) Signed-off-by: David Chen <530634352@qq.com> --- tests/entrypoints/openai_api/test_image_server.py | 10 ++++++---- .../test_async_omni_diffusion_config.py | 14 ++++++++++++++ vllm_omni/engine/async_omni_engine.py | 11 +++++++++++ 3 files changed, 31 insertions(+), 4 deletions(-) diff --git a/tests/entrypoints/openai_api/test_image_server.py b/tests/entrypoints/openai_api/test_image_server.py index 4b38692da..cd1e83393 100644 --- a/tests/entrypoints/openai_api/test_image_server.py +++ b/tests/entrypoints/openai_api/test_image_server.py @@ -177,7 +177,7 @@ def test_client(mock_async_diffusion): [BaseModelPath(name="Qwen/Qwen-Image", model_path="Qwen/Qwen-Image")] ) app.state.args = Namespace( - default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5}}', + default_sampling_params='{"0": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}', max_generated_image_size=1024 * 1792, ) @@ -200,7 +200,7 @@ def async_omni_test_client(): SimpleNamespace(stage_type="diffusion"), ] app.state.args = Namespace( - default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}', + default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}', max_generated_image_size=1048576, # 1024*1024 to support resolution tests ) return TestClient(app) @@ -222,7 +222,7 @@ def async_omni_rgba_test_client(): SimpleNamespace(stage_type="diffusion"), ] app.state.args = Namespace( - default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}', + default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}', max_generated_image_size=1048576, ) return TestClient(app) @@ -244,7 +244,7 @@ def async_omni_stage_configs_only_client(): # Intentionally do not populate app.state.stage_configs. Refactored # AsyncOmni exposes stage_configs on the engine instance. app.state.args = Namespace( - default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5}}', + default_sampling_params='{"1": {"num_inference_steps":4, "guidance_scale":7.5, "generator_device":"cpu"}}', max_generated_image_size=1024 * 1792, ) return TestClient(app) @@ -960,6 +960,7 @@ def test_image_edit_parameter_default(async_omni_test_client): assert captured_sampling_params.num_outputs_per_prompt == 1 assert captured_sampling_params.num_inference_steps == 4 assert captured_sampling_params.guidance_scale == 7.5 + assert captured_sampling_params.generator_device == "cpu" # Test that a size exceeding max_generated_image_size returns 400 response = async_omni_test_client.post( @@ -993,6 +994,7 @@ def test_image_edit_parameter_default_single_stage(test_client): assert captured_sampling_params.num_outputs_per_prompt == 1 assert captured_sampling_params.num_inference_steps == 4 assert captured_sampling_params.guidance_scale == 7.5 + assert captured_sampling_params.generator_device == "cpu" # Size exceeding max_generated_image_size (1024*1792) returns 400 response = test_client.post( diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index a55eaf05b..7ed812826 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -69,6 +69,20 @@ def test_default_stage_config_propagates_ulysses_mode(): assert parallel_config.ulysses_mode == "advanced_uaa" +def test_default_stage_config_includes_default_sampling_params(): + """Ensure default sampling params survive the default diffusion-stage builder.""" + stage_cfg = AsyncOmniEngine._create_default_diffusion_stage_cfg( + { + "default_sampling_params": '{"0": {"generator_device":"cpu", "guidance_scale":7.5}}', + } + )[0] + + assert stage_cfg["default_sampling_params"] == { + "generator_device": "cpu", + "guidance_scale": 7.5, + } + + def test_serve_cli_accepts_ulysses_mode(): """Ensure diffusion serve CLI exposes ulysses_mode and wires it to parallel_config.""" parser = FlexibleArgumentParser() diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 23a85e9f5..01c11cb96 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1138,6 +1138,16 @@ class AsyncOmniEngine: # We temporally create a default config for diffusion stage. # In the future, we should merge the default config with the user-provided config. normalized_kwargs = dict(kwargs) + default_sampling_params = normalized_kwargs.get("default_sampling_params") + if isinstance(default_sampling_params, str): + try: + default_sampling_params = json.loads(default_sampling_params) + except json.JSONDecodeError: + logger.warning("Invalid default_sampling_params JSON, ignoring stage defaults.") + default_sampling_params = None + if not isinstance(default_sampling_params, dict): + default_sampling_params = None + stage_default_sampling_params = default_sampling_params.get("0", {}) if default_sampling_params else {} # TODO: hack, convert dtype to string to avoid non-premitive omegaconf create error. if "dtype" in normalized_kwargs and not isinstance(normalized_kwargs["dtype"], str): @@ -1234,6 +1244,7 @@ class AsyncOmniEngine: "devices": devices, }, "engine_args": stage_engine_args, + "default_sampling_params": stage_default_sampling_params, "final_output": True, "final_output_type": "image", } From d46397809852cd5599e659a63133cced30e549e5 Mon Sep 17 00:00:00 2001 From: Alex Brooks Date: Fri, 17 Apr 2026 02:02:30 -0600 Subject: [PATCH 36/97] [Model] Support Flux1 Schnell (#2528) Signed-off-by: Alex Brooks --- docs/models/supported_models.md | 1 + docs/user_guide/diffusion_features.md | 1 + .../offline_inference/text_to_image/README.md | 3 +- tests/e2e/offline_inference/test_flux.py | 37 +++++++++++++++++++ .../diffusion/models/flux/flux_transformer.py | 4 +- .../diffusion/models/flux/pipeline_flux.py | 7 +++- 6 files changed, 49 insertions(+), 4 deletions(-) create mode 100644 tests/e2e/offline_inference/test_flux.py diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index 0f9c8fff6..b4976b09c 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -47,6 +47,7 @@ th { | `Flux2KleinPipeline` | FLUX.2-klein | `black-forest-labs/FLUX.2-klein-4B`, `black-forest-labs/FLUX.2-klein-9B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `FluxKontextPipeline` | FLUX.1-Kontext-dev | `black-forest-labs/FLUX.1-Kontext-dev` | ✅︎ | ✅︎ | | | | `FluxPipeline` | FLUX.1-dev | `black-forest-labs/FLUX.1-dev` | ✅︎ | ✅︎ | | ✅︎ | +| `FluxPipeline` | FLUX.1-schnell | `black-forest-labs/FLUX.1-schnell` | ✅︎ | ✅︎ | | ✅︎ | | `OmniGen2Pipeline` | OmniGen2 | `OmniGen2/OmniGen2` | ✅︎ | ✅︎ | | ✅︎ | | `StableAudioPipeline` | Stable-Audio-Open | `stabilityai/stable-audio-open-1.0` | ✅︎ | ✅︎ | | ✅︎ | | `Qwen3TTSForConditionalGeneration` | Qwen3-TTS-12Hz-1.7B-CustomVoice | `Qwen/Qwen3-TTS-12Hz-1.7B-CustomVoice` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | diff --git a/docs/user_guide/diffusion_features.md b/docs/user_guide/diffusion_features.md index 45953b852..4e7003cce 100644 --- a/docs/user_guide/diffusion_features.md +++ b/docs/user_guide/diffusion_features.md @@ -109,6 +109,7 @@ The following tables show which models support each feature: |-------|:----------:|:-----------:|:---------------------:|:--------------:|:-----------------:|:------:|:------------------------:|:--------------------:|:--------------:|:----------------:| | **Bagel** | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | ❌ | | **FLUX.1-dev** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | +| **FLUX.1-schnell** | ❌ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.2-klein** | ✅ | ✅ | ✅ | ✅ | ✅ | ✅ | ❌ | ❌ | ✅ | ❌ | | **FLUX.1-Kontext-dev** | ❌ | ❌ | ❌ | ❌ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | | **FLUX.2-dev** | ✅ | ✅ | ❌ | ✅ | ✅ | ✅ | ❌ | ❌ | ❌ | ❌ | diff --git a/examples/offline_inference/text_to_image/README.md b/examples/offline_inference/text_to_image/README.md index cc295e827..c71773972 100644 --- a/examples/offline_inference/text_to_image/README.md +++ b/examples/offline_inference/text_to_image/README.md @@ -29,7 +29,8 @@ This folder provides several entrypoints for experimenting with text-to-image di | `AIDC-AI/Ovis-Image-7B` | 1024 x 1024 | 71.8 | 17.1 | | `OmniGen2/OmniGen2` | 1024 x 1024 | 20.1 | 14.7 | | `stabilityai/stable-diffusion-3.5-medium` | 1024 x 1024 | 20.1 | 15.6 | -| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 77.6 | 31.4 | +| `black-forest-labs/FLUX.1-dev` | 1024 x 1024 | 33.9 | 31.4 | +| `black-forest-labs/FLUX.1-schnell` | 1024 x 1024 | 33.9 | 31.4 | | `black-forest-labs/FLUX.2-klein-4B` | 1024 x 1024 | 72.7 | 14.9 | | `black-forest-labs/FLUX.2-klein-9B` | 1024 x 1024 | 37.1 | 32.3 | | `black-forest-labs/FLUX.2-dev` | 1024 x 1024 | 65.7 | >80 (CPU offload required) | diff --git a/tests/e2e/offline_inference/test_flux.py b/tests/e2e/offline_inference/test_flux.py new file mode 100644 index 000000000..02c6787be --- /dev/null +++ b/tests/e2e/offline_inference/test_flux.py @@ -0,0 +1,37 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Tests for Flux1 Schnell.""" + +import pytest +from PIL import Image + +from vllm_omni.entrypoints.omni import Omni +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +MODEL = "black-forest-labs/FLUX.1-schnell" + + +@pytest.mark.core_model +@pytest.mark.diffusion +def test_flux_schnell_text_to_image(): + """Test FLUX.1-schnell text-to-image generation.""" + omni = Omni(model=MODEL) + + omni_outputs = list( + omni.generate( + prompts=["A photo of a cat sitting on a laptop"], + sampling_params_list=OmniDiffusionSamplingParams( + height=512, + width=512, + num_inference_steps=2, + seed=42, + ), + ) + ) + + assert len(omni_outputs) > 0 + images = omni_outputs[0].images + assert len(images) == 1 + assert isinstance(images[0], Image.Image) + assert images[0].size == (512, 512) diff --git a/vllm_omni/diffusion/models/flux/flux_transformer.py b/vllm_omni/diffusion/models/flux/flux_transformer.py index 297c62675..dd88ee76c 100644 --- a/vllm_omni/diffusion/models/flux/flux_transformer.py +++ b/vllm_omni/diffusion/models/flux/flux_transformer.py @@ -526,10 +526,10 @@ class FluxTransformer2DModel(nn.Module): def __init__( self, - od_config: OmniDiffusionConfig = None, + od_config: OmniDiffusionConfig | None = None, patch_size: int = 1, in_channels: int = 64, - out_channels: int = None, + out_channels: int | None = None, num_layers: int = 19, num_single_layers: int = 38, attention_head_dim: int = 128, diff --git a/vllm_omni/diffusion/models/flux/pipeline_flux.py b/vllm_omni/diffusion/models/flux/pipeline_flux.py index 6f43e8dbb..70d572d9a 100644 --- a/vllm_omni/diffusion/models/flux/pipeline_flux.py +++ b/vllm_omni/diffusion/models/flux/pipeline_flux.py @@ -30,6 +30,7 @@ from vllm_omni.diffusion.models.flux.flux_pipeline_mixin import FluxPipelineMixi from vllm_omni.diffusion.models.t5_encoder import T5EncoderModel from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import DiffusionPipelineProfilerMixin from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.diffusion.utils.tf_utils import get_transformer_config_kwargs from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific logger = logging.getLogger(__name__) @@ -106,7 +107,11 @@ class FluxPipeline(nn.Module, FluxPipelineMixin, CFGParallelMixin, DiffusionPipe self.vae = AutoencoderKL.from_pretrained(model, subfolder="vae", local_files_only=local_files_only).to( self.device ) - self.transformer = FluxTransformer2DModel(od_config=od_config, quant_config=od_config.quantization_config) + + transformer_kwargs = get_transformer_config_kwargs(od_config.tf_model_config, FluxTransformer2DModel) + self.transformer = FluxTransformer2DModel( + **transformer_kwargs, od_config=od_config, quant_config=od_config.quantization_config + ) self.tokenizer = CLIPTokenizer.from_pretrained(model, subfolder="tokenizer", local_files_only=local_files_only) self.tokenizer_2 = T5TokenizerFast.from_pretrained( From bbd6a443ed6fad62c334f34f3e226add89c13c40 Mon Sep 17 00:00:00 2001 From: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com> Date: Fri, 17 Apr 2026 16:13:00 +0800 Subject: [PATCH 37/97] [Core] Refactor CFG companion tracker and use in Orchestrator (#2623) Signed-off-by: yinpe <11810305@mail.sustech.edu.cn> --- docs/api/README.md | 2 +- tests/engine/test_cfg_companion_tracker.py | 82 ++++++ .../test_orchestrator_kv_sender_info.py | 5 +- .../entrypoints/test_cfg_companion_tracker.py | 114 --------- vllm_omni/engine/async_omni_engine.py | 1 - vllm_omni/engine/cfg_companion_tracker.py | 125 ++++++++++ vllm_omni/engine/orchestrator.py | 108 ++------ .../entrypoints/cfg_companion_tracker.py | 233 ------------------ 8 files changed, 235 insertions(+), 435 deletions(-) create mode 100644 tests/engine/test_cfg_companion_tracker.py delete mode 100644 tests/entrypoints/test_cfg_companion_tracker.py create mode 100644 vllm_omni/engine/cfg_companion_tracker.py delete mode 100644 vllm_omni/entrypoints/cfg_companion_tracker.py diff --git a/docs/api/README.md b/docs/api/README.md index f65cbb525..0147f19e1 100644 --- a/docs/api/README.md +++ b/docs/api/README.md @@ -5,7 +5,7 @@ Main entry points for vLLM-Omni inference and serving. - [vllm_omni.entrypoints.async_omni.AsyncOmni][] -- [vllm_omni.entrypoints.cfg_companion_tracker.CfgCompanionTracker][] +- [vllm_omni.engine.cfg_companion_tracker.CfgCompanionTracker][] - [vllm_omni.entrypoints.cli.benchmark.base.OmniBenchmarkSubcommandBase][] - [vllm_omni.entrypoints.cli.benchmark.main.OmniBenchmarkSubcommand][] - [vllm_omni.entrypoints.cli.benchmark.serve.OmniBenchmarkServingSubcommand][] diff --git a/tests/engine/test_cfg_companion_tracker.py b/tests/engine/test_cfg_companion_tracker.py new file mode 100644 index 000000000..f856a38c3 --- /dev/null +++ b/tests/engine/test_cfg_companion_tracker.py @@ -0,0 +1,82 @@ +import pytest + +from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def test_register_companion_and_cleanup(): + tracker = CfgCompanionTracker() + + tracker.register_companion("req1", "cfg_text", "req1__cfg_text") + tracker.register_companion("req1", "cfg_img", "req1__cfg_img") + + assert tracker.is_companion("req1__cfg_text") + assert tracker.get_companion_request_ids("req1") == { + "cfg_text": "req1__cfg_text", + "cfg_img": "req1__cfg_img", + } + + removed = tracker.cleanup_parent("req1") + + assert sorted(removed) == ["req1__cfg_img", "req1__cfg_text"] + assert not tracker.is_companion("req1__cfg_text") + assert tracker.get_companion_request_ids("req1") == {} + + +def test_attach_cfg_request_ids_clones_diffusion_params(): + tracker = CfgCompanionTracker() + tracker.register_companion("req1", "cfg_text", "req1__cfg_text") + + params = OmniDiffusionSamplingParams() + updated = tracker.attach_cfg_request_ids("req1", params) + + assert updated is not params + assert params.cfg_kv_request_ids is None + assert updated.cfg_kv_request_ids == {"cfg_text": "req1__cfg_text"} + + +def test_abort_parent_expands_to_companions_and_cleans_up_deferred_parent(): + tracker = CfgCompanionTracker() + tracker.register_companion("req1", "cfg_text", "req1__cfg_text") + tracker.defer_parent("req1", {"out": 1}, stage_id=0) + + aborted = tracker.abort_parents(["req1"]) + + assert sorted(aborted) == ["req1", "req1__cfg_text"] + assert not tracker.is_companion("req1__cfg_text") + assert tracker.pop_pending_parent("req1") is None + + +def test_abort_companion_does_not_expand_to_parent(): + tracker = CfgCompanionTracker() + tracker.register_companion("req1", "cfg_text", "req1__cfg_text") + + aborted = tracker.abort_parents(["req1__cfg_text"]) + + assert aborted == ["req1__cfg_text"] + + +def test_companion_completion_flushes_deferred_parent(): + tracker = CfgCompanionTracker() + tracker.register_companion("req1", "cfg_text", "req1__cfg_text") + tracker.defer_parent("req1", {"out": 123}, stage_id=0) + + assert not tracker.all_companions_done("req1") + assert tracker.on_companion_completed("req1__cfg_text") == "req1" + assert tracker.all_companions_done("req1") + + popped = tracker.pop_pending_parent("req1") + assert popped is not None + assert popped["engine_outputs"] == {"out": 123} + assert popped["stage_id"] == 0 + + +def test_companion_completion_without_registered_parent_asserts(): + tracker = CfgCompanionTracker() + tracker._companion_ids.add("req1__cfg_text") + tracker._companion_to_parent["req1__cfg_text"] = "req1" + + with pytest.raises(AssertionError, match="completed before parent req1 was registered"): + tracker.on_companion_completed("req1__cfg_text") diff --git a/tests/engine/test_orchestrator_kv_sender_info.py b/tests/engine/test_orchestrator_kv_sender_info.py index 94da4ce71..7e3fe0906 100644 --- a/tests/engine/test_orchestrator_kv_sender_info.py +++ b/tests/engine/test_orchestrator_kv_sender_info.py @@ -4,6 +4,7 @@ from types import SimpleNamespace import pytest from vllm import SamplingParams +from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.orchestrator import Orchestrator, OrchestratorRequestState from vllm_omni.engine.stage_engine_core_client import StageEngineCoreClient from vllm_omni.inputs.data import OmniDiffusionSamplingParams @@ -130,7 +131,7 @@ def test_forward_to_diffusion_attaches_kv_sender_info(): orchestrator.num_stages = 2 orchestrator.stage_clients = [sender_stage, diffusion_stage] - orchestrator._companion_map = {} + orchestrator._cfg_tracker = CfgCompanionTracker() orchestrator.stage_vllm_configs = [None, None] orchestrator.output_processors = [None, None] @@ -161,7 +162,7 @@ def test_forward_to_diffusion_uses_engine_input_source_for_kv_sender_info(): orchestrator.num_stages = 3 orchestrator.stage_clients = [source_stage, previous_stage, diffusion_stage] - orchestrator._companion_map = {} + orchestrator._cfg_tracker = CfgCompanionTracker() orchestrator.stage_vllm_configs = [None, None, None] orchestrator.output_processors = [None, None, None] diff --git a/tests/entrypoints/test_cfg_companion_tracker.py b/tests/entrypoints/test_cfg_companion_tracker.py deleted file mode 100644 index 941ead41f..000000000 --- a/tests/entrypoints/test_cfg_companion_tracker.py +++ /dev/null @@ -1,114 +0,0 @@ -import time -from types import SimpleNamespace - -import pytest - -from vllm_omni.entrypoints.cfg_companion_tracker import CfgCompanionTracker - -pytestmark = [pytest.mark.core_model, pytest.mark.cpu] - - -def dummy_expand_func(prompt, sp0): - if prompt == "expand_me": - return [SimpleNamespace(prompt={"prompt": "neg"}, role="cfg_text", request_id_suffix="__cfg_text")] - return [] - - -@pytest.fixture -def tracker(): - sp0 = SimpleNamespace() - return CfgCompanionTracker(prompt_expand_func=dummy_expand_func, stage0_sampling_params=sp0, timeout_s=0.1) - - -def test_companion_tracker_initialization(tracker): - assert not tracker.is_active - assert tracker.num_companions == 0 - - -def test_expand_prompts_registers_companions(tracker): - request_id_to_prompt = {"req1": "expand_me", "req2": "do_not_expand"} - - pairs = tracker.expand_prompts(request_id_to_prompt) - - assert len(pairs) == 1 - companion_id, prompt = pairs[0] - assert companion_id == "req1__cfg_text" - assert prompt == {"prompt": "neg"} - - assert tracker.is_active - assert tracker.num_companions == 1 - assert tracker.is_companion("req1__cfg_text") - assert not tracker.is_companion("req2__cfg_text") - assert tracker.has_companions("req1") - assert not tracker.has_companions("req2") - - comp_map = tracker.get_companion_request_ids("req1") - assert comp_map == {"cfg_text": "req1__cfg_text"} - - -def test_companion_lifecycle_success(tracker): - request_id_to_prompt = {"req1": "expand_me"} - tracker.expand_prompts(request_id_to_prompt) - - # Defer parent - engine_outputs = {"out": 123} - tracker.defer_parent("req1", engine_outputs, stage_id=0) - - # Initially not done - assert not tracker.all_companions_done("req1") - - # Companion completes - parent_id = tracker.on_companion_completed("req1__cfg_text") - - # Parent should be returned since all companions are done and it is pending - assert parent_id == "req1" - assert tracker.all_companions_done("req1") - - # Pop pending parent - popped = tracker.pop_pending_parent("req1") - assert popped is not None - assert popped["engine_outputs"] == engine_outputs - assert popped["stage_id"] == 0 - - -def test_companion_lifecycle_failure(tracker): - request_id_to_prompt = {"req1": "expand_me"} - tracker.expand_prompts(request_id_to_prompt) - - tracker.defer_parent("req1", {"out": 123}, stage_id=0) - - # Companion fails - parent_id, aborted = tracker.on_companion_error("req1__cfg_text") - - assert parent_id == "req1" - assert aborted is True - assert tracker.is_parent_failed("req1") - - # Parent should be removed from pending list - assert tracker.pop_pending_parent("req1") is None - - # Consume failure - tracker.consume_parent_failure("req1") - assert not tracker.is_parent_failed("req1") - - -def test_companion_lifecycle_timeout(tracker): - request_id_to_prompt = {"req1": "expand_me"} - tracker.expand_prompts(request_id_to_prompt) - - tracker.defer_parent("req1", {"out": 123}, stage_id=0) - - # Initially no timeouts - timeouts = tracker.check_timeouts() - assert len(timeouts) == 0 - - # Wait for timeout - time.sleep(0.15) - - # Check timeouts again - timeouts = tracker.check_timeouts() - assert len(timeouts) == 1 - assert timeouts[0] == "req1" - - # Should be removed from pending - assert tracker.pop_pending_parent("req1") is None diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 01c11cb96..238bbdcdb 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1072,7 +1072,6 @@ class AsyncOmniEngine: params=companion_params, supported_tasks=self.supported_tasks, ) - request = _upgrade_to_omni_request(request, companion_prompt) request.external_req_id = cid self.output_processors[0].add_request( diff --git a/vllm_omni/engine/cfg_companion_tracker.py b/vllm_omni/engine/cfg_companion_tracker.py new file mode 100644 index 000000000..b9dfae833 --- /dev/null +++ b/vllm_omni/engine/cfg_companion_tracker.py @@ -0,0 +1,125 @@ +"""CFG companion request tracker for the Omni orchestrator. + +Encapsulates all bookkeeping for Classifier-Free Guidance companion +requests (parent/companion ID mapping, completion tracking, +deferred forwarding, and cleanup). +""" + +from __future__ import annotations + +import logging +from typing import Any + +from vllm_omni.inputs.data import OmniDiffusionSamplingParams + +logger = logging.getLogger(__name__) + + +class CfgCompanionTracker: + """Manages CFG companion request lifecycle in the orchestrator scheduling loop.""" + + def __init__(self) -> None: + self._companion_map: dict[str, dict[str, str]] = {} # parent -> {role: companion_id} + self._companion_ids: set[str] = set() + self._companion_to_parent: dict[str, str] = {} # companion -> parent + self._done: dict[str, set[str]] = {} # parent -> completed companion ids + self._pending_parents: dict[str, dict[str, Any]] = {} # parent -> deferred result + + def is_companion(self, req_id: str) -> bool: + return req_id in self._companion_ids + + def has_companions(self, parent_id: str) -> bool: + return parent_id in self._companion_map + + def all_companions_done(self, parent_id: str) -> bool: + role_map = self._companion_map.get(parent_id, {}) + done_set = self._done.get(parent_id, set()) + return all(cid in done_set for cid in role_map.values()) + + def get_companion_request_ids(self, parent_id: str) -> dict[str, str]: + """Return ``{role: companion_request_id}`` for a parent.""" + return self._companion_map.get(parent_id, {}) + + def register_parent(self, parent_id: str) -> None: + self._companion_map.setdefault(parent_id, {}) + self._done.setdefault(parent_id, set()) + + def register_companion(self, parent_id: str, role: str, companion_id: str) -> None: + self.register_parent(parent_id) + self._companion_map[parent_id][role] = companion_id + self._companion_ids.add(companion_id) + self._companion_to_parent[companion_id] = parent_id + + def attach_cfg_request_ids(self, parent_id: str, sampling_params: Any) -> Any: + cfg_ids = self.get_companion_request_ids(parent_id) + if not cfg_ids: + return sampling_params + + if isinstance(sampling_params, OmniDiffusionSamplingParams): + sampling_params = sampling_params.clone() + sampling_params.cfg_kv_request_ids = cfg_ids + logger.info( + "Attaching cfg_kv_request_ids=%s to request %s", + cfg_ids, + parent_id, + ) + return sampling_params + + def on_companion_completed(self, companion_id: str) -> str | None: + """Mark done. Returns parent_id only if parent is pending and all companions finished.""" + parent_id = self._companion_to_parent.get(companion_id) + if parent_id is None: + return None + done_set = self._done.get(parent_id) + assert done_set is not None, f"Companion {companion_id} completed before parent {parent_id} was registered" + if companion_id in done_set: + return None + done_set.add(companion_id) + logger.debug("CFG companion %s completed (parent=%s)", companion_id, parent_id) + if parent_id in self._pending_parents and self.all_companions_done(parent_id): + return parent_id + return None + + def defer_parent(self, parent_id: str, engine_outputs: Any, stage_id: int) -> None: + """Hold parent result while waiting for companions to finish.""" + # TODO: Add timeout/error recovery when the orchestrator grows a + # companion-failure path. Today deferred parents are released only when + # companions finish or the external layer aborts the request. + self._pending_parents[parent_id] = { + "engine_outputs": engine_outputs, + "stage_id": stage_id, + } + logger.debug("Parent %s deferred, waiting for CFG companions", parent_id) + + def pop_pending_parent(self, parent_id: str) -> dict[str, Any] | None: + return self._pending_parents.pop(parent_id, None) + + def cleanup_parent(self, parent_id: str) -> list[str]: + companion_ids = list(self._companion_map.pop(parent_id, {}).values()) + for companion_id in companion_ids: + self._companion_ids.discard(companion_id) + self._companion_to_parent.pop(companion_id, None) + self._done.pop(parent_id, None) + self._pending_parents.pop(parent_id, None) + return companion_ids + + def abort_parents(self, request_ids: list[str]) -> list[str]: + all_request_ids = list(request_ids) + seen = set(all_request_ids) + parents_to_cleanup: set[str] = set() + + for req_id in request_ids: + # The orchestrator calls this with parent request IDs. If a raw + # companion ID is passed here, keep it as a direct abort target and + # avoid tearing down parent tracking state implicitly. + if req_id not in self._companion_ids: + parents_to_cleanup.add(req_id) + + for parent_id in parents_to_cleanup: + companion_ids = self.cleanup_parent(parent_id) + for companion_id in companion_ids: + if companion_id not in seen: + seen.add(companion_id) + all_request_ids.append(companion_id) + + return all_request_ids diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 0fdab9c0d..1de4282fe 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -27,6 +27,7 @@ from vllm_omni.distributed.omni_connectors.adapter import compute_talker_prompt_ from vllm_omni.engine import ( OmniEngineCoreRequest, ) +from vllm_omni.engine.cfg_companion_tracker import CfgCompanionTracker from vllm_omni.engine.serialization import serialize_additional_information from vllm_omni.metrics.stats import StageRequestStats as StageRequestMetrics from vllm_omni.metrics.stats import StageStats @@ -138,11 +139,7 @@ class Orchestrator: self.request_states: dict[str, OrchestratorRequestState] = {} # CFG companion tracking - self._companion_map: dict[str, dict[str, str]] = {} - self._companion_to_parent: dict[str, str] = {} - self._companion_ids: set[str] = set() - self._companion_done: dict[str, set[str]] = {} - self._deferred_parents: dict[str, dict[str, Any]] = {} + self._cfg_tracker = CfgCompanionTracker() # Per-stage metrics accumulators. self._batch_seq: list[int] = [0] * self.num_stages @@ -334,7 +331,7 @@ class Orchestrator: # CFG companion handling: companions don't produce user-visible output # and don't forward to the next stage directly. - if finished and req_id in self._companion_ids: + if finished and self._cfg_tracker.is_companion(req_id): await self._handle_cfg_companion_ready(req_id) self.request_states.pop(req_id, None) return @@ -368,57 +365,34 @@ class Orchestrator: and not self.async_chunk and not self._next_stage_already_submitted(stage_id, req_state) ): - if req_id in self._companion_map and not self._all_companions_done(req_id): - self._deferred_parents[req_id] = { - "stage_id": stage_id, - "output": output, - } + if self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id): + self._cfg_tracker.defer_parent(req_id, output, stage_id) else: await self._forward_to_next_stage(req_id, stage_id, output, req_state) if finished and stage_id == req_state.final_stage_id: - self._cleanup_companion_state(req_id) + self._cfg_tracker.cleanup_parent(req_id) self.request_states.pop(req_id, None) - def _cleanup_companion_state(self, parent_id: str) -> None: - """Remove all companion tracking state for a completed parent.""" - role_map = self._companion_map.pop(parent_id, {}) - for cid in role_map.values(): - self._companion_ids.discard(cid) - self._companion_to_parent.pop(cid, None) - self._companion_done.pop(parent_id, None) - self._deferred_parents.pop(parent_id, None) - - def _all_companions_done(self, parent_id: str) -> bool: - """Check whether all CFG companions for a parent request have finished.""" - role_map = self._companion_map.get(parent_id, {}) - if not role_map: - return True - done_set = self._companion_done.get(parent_id, set()) - return all(cid in done_set for cid in role_map.values()) - def _next_stage_already_submitted(self, stage_id: int, req_state: OrchestratorRequestState) -> bool: return (stage_id + 1) in req_state.stage_submit_ts async def _handle_cfg_companion_ready(self, req_id: str) -> None: """Mark a CFG companion as done; if all companions are done, flush deferred parent.""" - parent_id = self._companion_to_parent.get(req_id) + parent_id = self._cfg_tracker.on_companion_completed(req_id) if parent_id is None: return - done_set = self._companion_done.setdefault(parent_id, set()) - if req_id in done_set: + deferred = self._cfg_tracker.pop_pending_parent(parent_id) + if deferred is None: return - done_set.add(req_id) - if parent_id in self._deferred_parents and self._all_companions_done(parent_id): - deferred = self._deferred_parents.pop(parent_id) - parent_state = self.request_states.get(parent_id) - if parent_state is not None and not self._next_stage_already_submitted(deferred["stage_id"], parent_state): - await self._forward_to_next_stage( - parent_id, - deferred["stage_id"], - deferred["output"], - parent_state, - ) + parent_state = self.request_states.get(parent_id) + if parent_state is not None and not self._next_stage_already_submitted(deferred["stage_id"], parent_state): + await self._forward_to_next_stage( + parent_id, + deferred["stage_id"], + deferred["engine_outputs"], + parent_state, + ) async def _handle_kv_ready_raw_outputs(self, stage_id: int, raw_outputs: EngineCoreOutputs) -> None: """Forward split requests once stage-0 KV is ready, not only when decode fully finishes.""" @@ -432,18 +406,15 @@ class Orchestrator: req_state = self.request_states.get(req_id) if req_state is None: continue - if req_id in self._companion_ids: + if self._cfg_tracker.is_companion(req_id): await self._handle_cfg_companion_ready(req_id) continue if stage_id >= req_state.final_stage_id: continue if self._next_stage_already_submitted(stage_id, req_state): continue - if req_id in self._companion_map and not self._all_companions_done(req_id): - self._deferred_parents[req_id] = { - "stage_id": stage_id, - "output": raw_output, - } + if self._cfg_tracker.has_companions(req_id) and not self._cfg_tracker.all_companions_done(req_id): + self._cfg_tracker.defer_parent(req_id, raw_output, stage_id) else: await self._forward_to_next_stage(req_id, stage_id, raw_output, req_state) @@ -548,20 +519,7 @@ class Orchestrator: else: diffusion_prompt = req_state.prompt - # Attach CFG companion KV request IDs so the diffusion model - # runner can fetch companion KV caches alongside the primary one. - cfg_ids = self._companion_map.get(req_id) - if cfg_ids: - from vllm_omni.inputs.data import OmniDiffusionSamplingParams - - if isinstance(params, OmniDiffusionSamplingParams): - params = copy.deepcopy(params) - params.cfg_kv_request_ids = cfg_ids - logger.info( - "[Orchestrator] Attaching cfg_kv_request_ids=%s to req %s", - cfg_ids, - req_id, - ) + params = self._cfg_tracker.attach_cfg_request_ids(req_id, params) source_stage_ids = list(getattr(next_client, "engine_input_source", None) or [stage_id]) kv_sender_info = self._build_kv_sender_info(sender_stage_ids=source_stage_ids) @@ -817,13 +775,7 @@ class Orchestrator: companion_prompt = msg["prompt"] sampling_params_list = msg["sampling_params_list"] - # Register companion mapping - if parent_id not in self._companion_map: - self._companion_map[parent_id] = {} - self._companion_map[parent_id][role] = companion_id - self._companion_ids.add(companion_id) - self._companion_to_parent[companion_id] = parent_id - self._companion_done.setdefault(parent_id, set()) + self._cfg_tracker.register_companion(parent_id, role, companion_id) companion_state = OrchestratorRequestState( request_id=companion_id, @@ -848,22 +800,10 @@ class Orchestrator: async def _handle_abort(self, msg: dict[str, Any]) -> None: """Handle an abort message from the main thread.""" request_ids = msg["request_ids"] - # Also abort any CFG companions for aborted parents - companion_ids_to_abort: list[str] = [] - for req_id in request_ids: - role_map = self._companion_map.pop(req_id, {}) - for cid in role_map.values(): - companion_ids_to_abort.append(cid) - self._companion_ids.discard(cid) - self._companion_to_parent.pop(cid, None) - self.request_states.pop(cid, None) - self._companion_done.pop(req_id, None) - self._deferred_parents.pop(req_id, None) - - all_ids_to_abort = list(request_ids) + companion_ids_to_abort + all_ids_to_abort = self._cfg_tracker.abort_parents(request_ids) for stage_id in range(self.num_stages): await self.stage_clients[stage_id].abort_requests_async(all_ids_to_abort) - for req_id in request_ids: + for req_id in all_ids_to_abort: self.request_states.pop(req_id, None) logger.info("[Orchestrator] Aborted request(s) %s", request_ids) diff --git a/vllm_omni/entrypoints/cfg_companion_tracker.py b/vllm_omni/entrypoints/cfg_companion_tracker.py deleted file mode 100644 index 9c2e835f0..000000000 --- a/vllm_omni/entrypoints/cfg_companion_tracker.py +++ /dev/null @@ -1,233 +0,0 @@ -"""CFG companion request tracker for the Omni orchestrator. - -Encapsulates all bookkeeping for Classifier-Free Guidance companion -requests (prompt expansion, parent/companion ID mapping, completion -tracking, deferred forwarding, failure propagation, and timeouts) -so that ``Omni._run_generation`` stays clean. -""" - -from __future__ import annotations - -import copy -import logging -import os -import time -from collections.abc import Callable, Sequence -from typing import Any - -from vllm_omni.distributed.omni_connectors.adapter import try_send_via_connector -from vllm_omni.inputs.data import OmniDiffusionSamplingParams, OmniSamplingParams - -logger = logging.getLogger(__name__) - - -class CfgCompanionTracker: - """Manages CFG companion request lifecycle in the orchestrator scheduling loop.""" - - def __init__( - self, - prompt_expand_func: Callable[..., Any] | None, - stage0_sampling_params: Any, - timeout_s: float | None = None, - ) -> None: - self._expand_func = prompt_expand_func - self._sp0 = stage0_sampling_params - self._timeout_s = ( - timeout_s if timeout_s is not None else float(os.environ.get("VLLM_CFG_PENDING_TIMEOUT_S", "120")) - ) - - self._companion_map: dict[str, dict[str, str]] = {} # parent -> {role: companion_id} - self._companion_ids: set[str] = set() - self._companion_to_parent: dict[str, str] = {} # companion -> parent - self._done: dict[str, set[str]] = {} # parent -> completed companion ids - self._pending_parents: dict[str, dict[str, Any]] = {} # parent -> deferred result - self._failed_parents: set[str] = set() - - @property - def is_active(self) -> bool: - return bool(self._companion_ids) - - @property - def num_companions(self) -> int: - return len(self._companion_ids) - - @property - def stage0_sampling_params(self) -> Any: - return self._sp0 - - def expand_prompts( - self, - request_id_to_prompt: dict[str, Any], - ) -> list[tuple[str, Any]]: - """Expand user prompts into ``(companion_id, prompt)`` pairs via model-specific func.""" - if not self._expand_func: - return [] - - pairs: list[tuple[str, Any]] = [] - for rid, prompt in request_id_to_prompt.items(): - expanded = self._expand_func(prompt, self._sp0) - if not expanded: - continue - role_map: dict[str, str] = {} - for ep in expanded: - cid = f"{rid}{ep.request_id_suffix}" - role_map[ep.role] = cid - self._companion_ids.add(cid) - self._companion_to_parent[cid] = rid - pairs.append((cid, ep.prompt)) - self._companion_map[rid] = role_map - self._done[rid] = set() - - logger.debug( - "CFG expansion: %d parent(s) -> %d companion(s)", - len(self._companion_map), - len(self._companion_ids), - ) - return pairs - - def is_companion(self, req_id: str) -> bool: - return req_id in self._companion_ids - - def has_companions(self, parent_id: str) -> bool: - return parent_id in self._companion_map - - def all_companions_done(self, parent_id: str) -> bool: - role_map = self._companion_map.get(parent_id, {}) - done_set = self._done.get(parent_id, set()) - return all(cid in done_set for cid in role_map.values()) - - def get_companion_request_ids(self, parent_id: str) -> dict[str, str]: - """Return ``{role: companion_request_id}`` for a parent.""" - return self._companion_map.get(parent_id, {}) - - def is_parent_failed(self, parent_id: str) -> bool: - return parent_id in self._failed_parents - - # -- Lifecycle events -- - - def on_companion_error(self, companion_id: str) -> tuple[str | None, bool]: - """Record failure. Returns ``(parent_id, parent_was_aborted)``.""" - parent_id = self._companion_to_parent.get(companion_id) - if parent_id is None: - return None, False - self._failed_parents.add(parent_id) - logger.error("CFG companion %s failed; marking parent %s as failed", companion_id, parent_id) - aborted = parent_id in self._pending_parents - if aborted: - self._pending_parents.pop(parent_id, None) - return parent_id, aborted - - def on_companion_completed(self, companion_id: str) -> str | None: - """Mark done. Returns parent_id only if parent is pending and all companions finished.""" - parent_id = self._companion_to_parent.get(companion_id) - if parent_id is None: - return None - self._done[parent_id].add(companion_id) - logger.debug("CFG companion %s completed (parent=%s)", companion_id, parent_id) - if parent_id in self._pending_parents and self.all_companions_done(parent_id): - return parent_id - return None - - def consume_parent_failure(self, parent_id: str) -> None: - self._failed_parents.discard(parent_id) - - # -- Deferred parent management -- - - def defer_parent(self, parent_id: str, engine_outputs: Any, stage_id: int) -> None: - """Hold parent result while waiting for companions to finish.""" - self._pending_parents[parent_id] = { - "engine_outputs": engine_outputs, - "stage_id": stage_id, - "pending_since": time.time(), - } - logger.debug("Parent %s deferred, waiting for CFG companions", parent_id) - - def pop_pending_parent(self, parent_id: str) -> dict[str, Any] | None: - return self._pending_parents.pop(parent_id, None) - - def check_timeouts(self) -> list[str]: - """Return and remove parent IDs that exceeded the pending timeout.""" - if not self._pending_parents: - return [] - now = time.time() - timed_out: list[str] = [] - for pid in list(self._pending_parents): - pending_since = self._pending_parents[pid].get("pending_since", now) - if now - pending_since > self._timeout_s: - self._pending_parents.pop(pid) - self._failed_parents.discard(pid) - timed_out.append(pid) - logger.error("Parent %s timed out waiting for CFG companions (>%.0fs)", pid, self._timeout_s) - return timed_out - - # -- Forward parent with CFG KV -- - - def forward_parent_with_cfg( - self, - req_id: str, - parent_result: dict[str, Any], - stage_list: Sequence[Any], - connectors: dict[tuple[str, str], Any], - sampling_params_list: Sequence[OmniSamplingParams], - request_id_to_prompt: dict[str, Any], - final_stage_id_to_prompt: dict[str, int], - metrics: Any, - remaining_by_stage: list[int], - ) -> bool: - """Forward a parent request to the next stage with CFG KV request IDs attached.""" - stage_id = parent_result["stage_id"] - next_stage_id = stage_id + 1 - if next_stage_id > final_stage_id_to_prompt.get(req_id, 0): - return True - - next_stage = stage_list[next_stage_id] - try: - with metrics.stage_postprocess_timer(stage_id, req_id): - next_inputs = next_stage.process_engine_inputs( - stage_list, - [request_id_to_prompt[req_id]], - source_outputs_override=parent_result["engine_outputs"], - ) - except Exception as e: - logger.exception( - "Process engine inputs error for req %s at stage %d: %s", - req_id, - next_stage_id, - e, - ) - return False - - sp_next = copy.deepcopy(sampling_params_list[next_stage_id]) - if isinstance(sp_next, OmniDiffusionSamplingParams): - sp_next.cfg_kv_request_ids = self.get_companion_request_ids(req_id) - logger.info( - "Attaching cfg_kv_request_ids=%s to request %s", - sp_next.cfg_kv_request_ids, - req_id, - ) - - connector_key = (str(stage_id), str(next_stage_id)) - connector = connectors.get(connector_key) - sent_via_connector = False - if connector: - sent_via_connector = try_send_via_connector( - connector=connector, - stage_id=stage_id, - next_stage_id=next_stage_id, - req_id=req_id, - next_inputs=next_inputs, - sampling_params=sp_next, - original_prompt=request_id_to_prompt[req_id], - next_stage_queue_submit_fn=stage_list[next_stage_id].submit, - metrics=metrics, - ) - - if not sent_via_connector: - raise RuntimeError( - f"Failed to send CFG request {req_id} to stage-{next_stage_id} via connector. " - "Configure a connector for this edge or inspect connector logs for details." - ) - - logger.debug("Forwarded CFG-enabled request %s to stage-%d", req_id, next_stage_id) - remaining_by_stage[next_stage_id] += 1 - return True From a3ecde98e0001481a0ed258850a7176789f60c0a Mon Sep 17 00:00:00 2001 From: wangyu <53896905+yenuo26@users.noreply.github.com> Date: Fri, 17 Apr 2026 16:33:46 +0800 Subject: [PATCH 38/97] [CI][Bugfix] Fix the error in generating the performance data table and add a fallback mechanism that prevents the result file from being generated when test case execution fails. (#2839) Signed-off-by: wangyu <410167048@qq.com> Signed-off-by: wangyu <53896905+yenuo26@users.noreply.github.com> --- .buildkite/test-nightly.yml | 64 ++++++++++++-- .../scripts/diffusion_result_template.json | 86 +++++++++++++++++++ .../perf/scripts/result_omni_template.json | 55 ++++++++++++ tests/dfx/perf/scripts/run_benchmark.py | 17 +++- .../perf/scripts/run_diffusion_benchmark.py | 12 ++- tests/test_generate_nightly_perf_excel.py | 71 +++++++++++++++ tests/test_generate_nightly_perf_html.py | 54 ++++++++++++ tools/nightly/generate_nightly_perf_excel.py | 16 ++++ tools/nightly/generate_nightly_perf_html.py | 50 +++++++---- 9 files changed, 394 insertions(+), 31 deletions(-) create mode 100644 tests/dfx/perf/scripts/diffusion_result_template.json create mode 100644 tests/dfx/perf/scripts/result_omni_template.json create mode 100644 tests/test_generate_nightly_perf_excel.py create mode 100644 tests/test_generate_nightly_perf_html.py diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index ac43b597d..02d8cced4 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -110,8 +110,12 @@ steps: timeout_in_minutes: 180 commands: - export BENCHMARK_DIR=tests/dfx/perf/results - - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_omni.json - - buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_omni.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT agents: queue: "mithril-h100-pool" plugins: @@ -177,8 +181,12 @@ steps: commands: - export BENCHMARK_DIR=tests/dfx/perf/results - export VLLM_ALLOW_LONG_MAX_MODEL_LEN="1" - - pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json - - buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + - | + set +e + pytest -s -v tests/dfx/perf/scripts/run_benchmark.py --test-config-file tests/dfx/perf/tests/test_tts.json + EXIT=$$? + buildkite-agent artifact upload "tests/dfx/perf/results/*.json" + exit $$EXIT agents: queue: "mithril-h100-pool" plugins: @@ -398,6 +406,46 @@ steps: path: /mnt/hf-cache type: DirectoryOrCreate + - label: ":full_moon: Diffusion X2I(&A&T) · Accuracy Test" + timeout_in_minutes: 180 + commands: + - pytest -s -v tests/e2e/accuracy/test_qwen_image*.py --run-level advanced_model + agents: + queue: "mithril-h100-pool" + plugins: + - kubernetes: + podSpec: + containers: + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT + resources: + limits: + nvidia.com/gpu: 1 + volumeMounts: + - name: devshm + mountPath: /dev/shm + - name: hf-cache + mountPath: /root/.cache/huggingface + env: + - name: HF_HOME + value: /root/.cache/huggingface + - name: VLLM_HTTP_TIMEOUT_KEEP_ALIVE + value: "120" + - name: HF_TOKEN + valueFrom: + secretKeyRef: + name: hf-token-secret + key: token + nodeSelector: + node.kubernetes.io/instance-type: gpu-h100-sxm + volumes: + - name: devshm + emptyDir: + medium: Memory + - name: hf-cache + hostPath: + path: /mnt/hf-cache + type: DirectoryOrCreate + - label: ":full_moon: Diffusion X2I(&A&T) · Perf Test" key: nightly-diffusion-x2iat-performance timeout_in_minutes: 180 @@ -417,11 +465,9 @@ steps: EXIT3=$$? pytest -s -v tests/dfx/perf/scripts/run_diffusion_benchmark.py --test-config-file tests/dfx/perf/tests/test_qwen_image_layered_vllm_omni.json EXIT4=$$? - if [ $$EXIT1 -eq 0 ] || [ $$EXIT2 -eq 0 ] || [ $$EXIT3 -eq 0 ] || [ $$EXIT4 -eq 0 ]; then - buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" - buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" - fi - exit $$((EXIT1 | EXIT2 | EXIT3)) + buildkite-agent artifact upload "tests/dfx/perf/results/diffusion_result_*.json" + buildkite-agent artifact upload "tests/dfx/perf/results/logs/*.log" + exit $$((EXIT1 | EXIT2 | EXIT3 | EXIT4)) agents: queue: "mithril-h100-pool" plugins: diff --git a/tests/dfx/perf/scripts/diffusion_result_template.json b/tests/dfx/perf/scripts/diffusion_result_template.json new file mode 100644 index 000000000..86bdf1bc7 --- /dev/null +++ b/tests/dfx/perf/scripts/diffusion_result_template.json @@ -0,0 +1,86 @@ +[ + { + "test_name": null, + "backend": null, + "timestamp": null, + "server_params": { + "model": null, + "serve_args": { + "enable-diffusion-pipeline-profiler": false + } + }, + "benchmark_params": { + "name": null, + "dataset": null, + "task": null, + "width": 0, + "height": 0, + "num-inference-steps": 0, + "num-prompts": 0, + "max-concurrency": 0, + "num-input-images": 0, + "enable-negative-prompt": false, + "baseline": { + "throughput_qps": 0, + "latency_mean": 0, + "peak_memory_mb_max": 0, + "peak_memory_mb_mean": 0 + } + }, + "result": { + "duration": 0, + "completed_requests": 0, + "failed_requests": 0, + "throughput_qps": 0, + "latency_mean": 0, + "latency_median": 0, + "latency_p99": 0, + "latency_p95": 0, + "latency_p50": 0, + "peak_memory_mb_max": 0, + "peak_memory_mb_mean": 0, + "peak_memory_mb_median": 0, + "stage_durations_mean": {}, + "stage_durations_p50": {}, + "stage_durations_p99": {}, + "backend": null, + "model": null, + "dataset": null, + "task": null + }, + "log_file": null, + "Model": null, + "Framework": null, + "Hardware": null, + "Deployment": null, + "Task": null, + "Dataset": null, + "resolution": null, + "Parallelism": null, + "max_concurrency": 0, + "Cache": null, + "Quantization": null, + "offload": null, + "compile": null, + "Attn_backend": null, + "num_inference_steps": 0, + "completed": 0, + "failed": 0, + "throughput_qps": 0, + "latency_mean": 0, + "latency_median": 0, + "latency_p99": 0, + "latency_p95": 0, + "latency_p50": 0, + "peak_memory_mb_max": 0, + "peak_memory_mb_mean": 0, + "peak_memory_mb_median": 0, + "stage_durations_mean": {}, + "stage_durations_p50": {}, + "stage_durations_p99": {}, + "commit_sha": null, + "build_id": null, + "build_url": null, + "source_file": null + } +] diff --git a/tests/dfx/perf/scripts/result_omni_template.json b/tests/dfx/perf/scripts/result_omni_template.json new file mode 100644 index 000000000..1d6132140 --- /dev/null +++ b/tests/dfx/perf/scripts/result_omni_template.json @@ -0,0 +1,55 @@ +{ + "date": null, + "endpoint_type": null, + "backend": null, + "label": null, + "model_id": null, + "tokenizer_id": null, + "num_prompts": 0, + "request_rate": null, + "burstiness": 0, + "max_concurrency": 0, + "duration": 0, + "completed": 0, + "failed": 0, + "total_input_tokens": 0, + "total_output_tokens": 0, + "request_throughput": 0, + "request_goodput": null, + "output_throughput": 0, + "total_token_throughput": 0, + "total_audio_duration_s": 0, + "total_audio_frames": 0, + "audio_throughput": 0, + "max_output_tokens_per_s": 0, + "max_concurrent_requests": 0, + "rtfx": 0, + "mean_ttft_ms": 0, + "median_ttft_ms": 0, + "p99_ttft_ms": 0, + "mean_tpot_ms": 0, + "median_tpot_ms": 0, + "p99_tpot_ms": 0, + "mean_itl_ms": 0, + "median_itl_ms": 0, + "p99_itl_ms": 0, + "mean_e2el_ms": 0, + "median_e2el_ms": 0, + "p99_e2el_ms": 0, + "mean_audio_rtf": 0, + "median_audio_rtf": 0, + "p99_audio_rtf": 0, + "mean_audio_ttfp_ms": 0, + "median_audio_ttfp_ms": 0, + "p99_audio_ttfp_ms": 0, + "mean_audio_duration_s": 0, + "median_audio_duration_s": 0, + "p99_audio_duration_s": 0, + "baseline": { + "mean_ttft_ms": 0, + "mean_audio_ttfp_ms": 0, + "mean_audio_rtf": 0 + }, + "random_input_len": 0, + "random_output_len": 0 +} diff --git a/tests/dfx/perf/scripts/run_benchmark.py b/tests/dfx/perf/scripts/run_benchmark.py index d5ef1b49e..bea46f684 100644 --- a/tests/dfx/perf/scripts/run_benchmark.py +++ b/tests/dfx/perf/scripts/run_benchmark.py @@ -45,6 +45,8 @@ if CONFIG_FILE_PATH is None: CONFIG_FILE_PATH = _DEFAULT_CONFIG_FILE BENCHMARK_CONFIGS = load_configs(CONFIG_FILE_PATH) +OMNI_RESULT_TEMPLATE_PATH = Path(__file__).parent / "result_omni_template.json" + STAGE_CONFIGS_DIR = Path(__file__).parent.parent / "stage_configs" test_params = create_unique_server_params(BENCHMARK_CONFIGS, STAGE_CONFIGS_DIR) @@ -64,7 +66,7 @@ def omni_server(request): print(f"Starting OmniServer with test: {test_name}, model: {model}") - server_args = ["--stage-init-timeout", "300", "--init-timeout", "900"] + server_args = ["--stage-init-timeout", "600", "--init-timeout", "900"] if stage_config_path: server_args = ["--stage-configs-path", stage_config_path] + server_args with OmniServer(model, server_args) as server: @@ -141,8 +143,17 @@ def run_benchmark( result_dir = "./" result_path = os.path.join(result_dir, result_filename) - with open(result_path, encoding="utf-8") as f: - result = json.load(f) + if not os.path.exists(result_path): + with open(OMNI_RESULT_TEMPLATE_PATH, encoding="utf-8") as f: + template_result: dict[str, Any] = json.load(f) + Path(result_path).parent.mkdir(parents=True, exist_ok=True) + with open(result_path, "w", encoding="utf-8") as f: + json.dump(template_result, f, ensure_ascii=False, indent=2) + print(f"Benchmark result file not generated, fallback to template: {result_path}") + result = template_result + else: + with open(result_path, encoding="utf-8") as f: + result = json.load(f) if baseline_config: result["baseline"] = _baseline_thresholds_for_step( diff --git a/tests/dfx/perf/scripts/run_diffusion_benchmark.py b/tests/dfx/perf/scripts/run_diffusion_benchmark.py index 8eeeec8df..23efa8bb0 100644 --- a/tests/dfx/perf/scripts/run_diffusion_benchmark.py +++ b/tests/dfx/perf/scripts/run_diffusion_benchmark.py @@ -52,6 +52,7 @@ BENCHMARK_SCRIPT = str( _SESSION_TIMESTAMP = datetime.now().strftime("%Y%m%d-%H%M%S") _RESULT_LOCK = threading.Lock() _BRANCHPOINT_COMMIT_SHA: str | None = None +DIFFUSION_RESULT_TEMPLATE_PATH = Path(__file__).parent / "diffusion_result_template.json" def _get_config_file_from_argv() -> str | None: @@ -563,10 +564,17 @@ def run_benchmark( if process.returncode != 0: tmp_result_file.unlink(missing_ok=True) - raise RuntimeError(f"Benchmark script exited with code {process.returncode}") + print(f"ERROR:Benchmark script exited with code {process.returncode}") if not tmp_result_file.exists(): - raise FileNotFoundError(f"Benchmark result file not found: {tmp_result_file}") + with open(DIFFUSION_RESULT_TEMPLATE_PATH, encoding="utf-8") as f: + template_payload = json.load(f) + # Template schema is fixed and owned by this repo: + # ``diffusion_result_template.json`` is a one-item list and metrics live at [0]["result"]. + template_metrics: dict[str, Any] = template_payload[0]["result"] + with open(tmp_result_file, "w", encoding="utf-8") as f: + json.dump(template_metrics, f, ensure_ascii=False, indent=2) + print(f"Benchmark result file not generated, fallback to template: {tmp_result_file}") try: with open(tmp_result_file, encoding="utf-8") as f: diff --git a/tests/test_generate_nightly_perf_excel.py b/tests/test_generate_nightly_perf_excel.py new file mode 100644 index 000000000..9b05d6de0 --- /dev/null +++ b/tests/test_generate_nightly_perf_excel.py @@ -0,0 +1,71 @@ +import importlib.util +import json +import sys +from pathlib import Path + +import pytest +from openpyxl import load_workbook + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _load_excel_module(): + repo_root = Path(__file__).resolve().parents[1] + module_path = repo_root / "tools" / "nightly" / "generate_nightly_perf_excel.py" + spec = importlib.util.spec_from_file_location("generate_nightly_perf_excel", module_path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def _cell_value_by_header(ws, header_name: str, row_idx: int = 2): + headers = [c.value for c in ws[1]] + col_idx = headers.index(header_name) + 1 + return ws.cell(row=row_idx, column=col_idx).value + + +def test_generate_excel_report_with_perf_templates(tmp_path: Path): + module = _load_excel_module() + repo_root = Path(__file__).resolve().parents[1] + perf_scripts_dir = repo_root / "tests" / "dfx" / "perf" / "scripts" + + omni_template_path = perf_scripts_dir / "result_omni_template.json" + diffusion_template_path = perf_scripts_dir / "diffusion_result_template.json" + + omni_record = json.loads(omni_template_path.read_text(encoding="utf-8")) + diffusion_records = json.loads(diffusion_template_path.read_text(encoding="utf-8")) + + input_dir = tmp_path / "input" + diffusion_input_dir = tmp_path / "diffusion_input" + input_dir.mkdir() + diffusion_input_dir.mkdir() + + # Keep file names compatible with parser conventions in generate_nightly_perf_excel.py + omni_result_file = input_dir / "result_test_perf_random_1_4_in2500_out900_20260415-185642.json" + diffusion_result_file = diffusion_input_dir / "diffusion_result_qwen_image_edit_20260415-193200.json" + omni_result_file.write_text(json.dumps(omni_record, ensure_ascii=False, indent=2), encoding="utf-8") + diffusion_result_file.write_text(json.dumps(diffusion_records, ensure_ascii=False, indent=2), encoding="utf-8") + + output_file = tmp_path / "nightly_perf.xlsx" + module.generate_excel_report( + input_dir=str(input_dir), + diffusion_input_dir=str(diffusion_input_dir), + output_file=str(output_file), + commit_sha="test_commit_sha", + build_id="test_build_id", + build_url="https://example.com/build/123", + ) + + assert output_file.exists() + + wb = load_workbook(output_file) + assert set(wb.sheetnames) >= {"omni_summary", "diffusion_summary", "omni_raw", "diffusion_raw"} + + ws_omni_raw = wb["omni_raw"] + baseline_cell = _cell_value_by_header(ws_omni_raw, "baseline") + assert baseline_cell == json.dumps(omni_record["baseline"], ensure_ascii=False, sort_keys=True) + + ws_omni_summary = wb["omni_summary"] + assert _cell_value_by_header(ws_omni_summary, "commit_sha") == "test_commit_sha" + assert _cell_value_by_header(ws_omni_summary, "build_id") == "test_build_id" diff --git a/tests/test_generate_nightly_perf_html.py b/tests/test_generate_nightly_perf_html.py new file mode 100644 index 000000000..4e77eb3ad --- /dev/null +++ b/tests/test_generate_nightly_perf_html.py @@ -0,0 +1,54 @@ +import importlib.util +import json +import sys +from pathlib import Path + +import pytest + +pytestmark = [pytest.mark.core_model, pytest.mark.cpu] + + +def _load_html_module(): + repo_root = Path(__file__).resolve().parents[1] + module_path = repo_root / "tools" / "nightly" / "generate_nightly_perf_html.py" + spec = importlib.util.spec_from_file_location("generate_nightly_perf_html", module_path) + module = importlib.util.module_from_spec(spec) + sys.modules[spec.name] = module + spec.loader.exec_module(module) + return module + + +def test_generate_html_report_with_perf_templates(tmp_path: Path): + module = _load_html_module() + repo_root = Path(__file__).resolve().parents[1] + perf_scripts_dir = repo_root / "tests" / "dfx" / "perf" / "scripts" + + omni_template_path = perf_scripts_dir / "result_omni_template.json" + diffusion_template_path = perf_scripts_dir / "diffusion_result_template.json" + + omni_record = json.loads(omni_template_path.read_text(encoding="utf-8")) + diffusion_records = json.loads(diffusion_template_path.read_text(encoding="utf-8")) + + input_dir = tmp_path / "input" + diffusion_input_dir = tmp_path / "diffusion_input" + input_dir.mkdir() + diffusion_input_dir.mkdir() + + omni_result_file = input_dir / "result_test_perf_random_1_4_in2500_out900_20260415-185642.json" + diffusion_result_file = diffusion_input_dir / "diffusion_result_qwen_image_edit_20260415-193200.json" + omni_result_file.write_text(json.dumps(omni_record, ensure_ascii=False, indent=2), encoding="utf-8") + diffusion_result_file.write_text(json.dumps(diffusion_records, ensure_ascii=False, indent=2), encoding="utf-8") + + output_file = tmp_path / "nightly_perf_v2.html" + module.generate_html_report( + input_dir=str(input_dir), + diffusion_input_dir=str(diffusion_input_dir), + output_file=str(output_file), + ) + + assert output_file.exists() + html = output_file.read_text(encoding="utf-8") + assert "Nightly Performance Report" in html + assert "Omni records 1" in html + assert f"Diffusion records {len(diffusion_records)}" in html + assert "const DIFF_DATA =" in html diff --git a/tools/nightly/generate_nightly_perf_excel.py b/tools/nightly/generate_nightly_perf_excel.py index 4bb778531..6ba1d1eef 100644 --- a/tools/nightly/generate_nightly_perf_excel.py +++ b/tools/nightly/generate_nightly_perf_excel.py @@ -665,6 +665,21 @@ def _to_float_if_numeric(value: Any) -> Any: return value +def _to_excel_compatible(value: Any) -> Any: + """Convert non-scalar objects to JSON string for openpyxl cell values.""" + if isinstance(value, (dict, list, tuple)): + try: + return json.dumps(value, ensure_ascii=False, sort_keys=True) + except (TypeError, ValueError): + return str(value) + if isinstance(value, set): + try: + return json.dumps(sorted(value), ensure_ascii=False) + except (TypeError, ValueError): + return str(value) + return value + + def _write_sheet( ws, columns: Sequence[str], @@ -680,6 +695,7 @@ def _write_sheet( v = record.get(col) if col in numeric_set: v = _to_float_if_numeric(v) + v = _to_excel_compatible(v) row_values.append(v) ws.append(row_values) diff --git a/tools/nightly/generate_nightly_perf_html.py b/tools/nightly/generate_nightly_perf_html.py index dd5ece790..a50a46255 100644 --- a/tools/nightly/generate_nightly_perf_html.py +++ b/tools/nightly/generate_nightly_perf_html.py @@ -17,7 +17,7 @@ from typing import Any LOGGER = logging.getLogger(__name__) _RESULT_JSON_PREFIX = "result_test_" -_DIFFUSION_JSON_PREFIX = "diffusion_perf_" +_DIFFUSION_JSON_PREFIXES = ("diffusion_perf_", "diffusion_result_") DEFAULT_INPUT_DIR = os.getenv("DEFAULT_INPUT_DIR") or "tests" DEFAULT_OUTPUT_DIR = os.getenv("DEFAULT_OUTPUT_DIR") or "tests" DEFAULT_DIFFUSION_INPUT_DIR = os.getenv("DIFFUSION_BENCHMARK_DIR") @@ -51,7 +51,7 @@ def _default_diffusion_input_dir(input_dir: str) -> str: return DEFAULT_DIFFUSION_INPUT_DIR if DEFAULT_DIFFUSION_INPUT_DIR else input_dir -def _load_json_file(path: str) -> dict[str, Any] | None: +def _load_json_file(path: str) -> dict[str, Any] | list[Any] | None: try: with open(path, encoding="utf-8") as f: data = json.load(f) @@ -59,8 +59,8 @@ def _load_json_file(path: str) -> dict[str, Any] | None: LOGGER.warning("failed to load json '%s': %s", path, exc) return None - if not isinstance(data, dict): - LOGGER.warning("json root in '%s' is not an object, skip", path) + if not isinstance(data, (dict, list)): + LOGGER.warning("json root in '%s' is not an object or list, skip", path) return None return data @@ -170,9 +170,10 @@ def _iter_omni_json_records(input_dir: str) -> Iterable[dict[str, Any]]: def _parse_diffusion_from_filename(filename: str) -> dict[str, Any]: name, ext = os.path.splitext(filename) - if ext != ".json" or not name.startswith(_DIFFUSION_JSON_PREFIX): + if ext != ".json" or not any(name.startswith(prefix) for prefix in _DIFFUSION_JSON_PREFIXES): return {} - core = name[len(_DIFFUSION_JSON_PREFIX) :] + matched_prefix = next(prefix for prefix in _DIFFUSION_JSON_PREFIXES if name.startswith(prefix)) + core = name[len(matched_prefix) :] parts = core.split("_") if len(parts) < 2: return {} @@ -195,7 +196,7 @@ def _iter_diffusion_json_records(input_dir: str) -> Iterable[dict[str, Any]]: return for entry in sorted(os.listdir(input_dir)): - if not entry.endswith(".json") or not entry.startswith(_DIFFUSION_JSON_PREFIX): + if not entry.endswith(".json") or not any(entry.startswith(prefix) for prefix in _DIFFUSION_JSON_PREFIXES): continue full_path = os.path.join(input_dir, entry) if not os.path.isfile(full_path): @@ -204,17 +205,32 @@ def _iter_diffusion_json_records(input_dir: str) -> Iterable[dict[str, Any]]: if data is None: continue - record: dict[str, Any] = dict(data) filename_meta = _parse_diffusion_from_filename(os.path.basename(full_path)) - if "date" not in record or not record.get("date"): - record["date"] = filename_meta.get("date") or datetime.now( - timezone.utc, - ).strftime("%Y%m%d-%H%M%S") - if "test_name" not in record or not record.get("test_name"): - if "test_name" in filename_meta: - record["test_name"] = filename_meta["test_name"] - record["source_file"] = os.path.basename(full_path) - yield record + if isinstance(data, dict): + records = [data] + elif isinstance(data, list): + records = [r for r in data if isinstance(r, dict)] + else: + records = [] + + if not records: + LOGGER.warning("diffusion json '%s' has no valid records, skip", full_path) + continue + + for record in records: + flat: dict[str, Any] = dict(record) + result = flat.pop("result", None) + if isinstance(result, dict): + flat.update(result) + if "date" not in flat or not flat.get("date"): + flat["date"] = filename_meta.get("date") or datetime.now( + timezone.utc, + ).strftime("%Y%m%d-%H%M%S") + if "test_name" not in flat or not flat.get("test_name"): + if "test_name" in filename_meta: + flat["test_name"] = filename_meta["test_name"] + flat["source_file"] = os.path.basename(full_path) + yield flat def _collect_omni_records(input_dir: str) -> list[dict[str, Any]]: From b88d3ce75ebd4ee52015e813b006e0080be5e463 Mon Sep 17 00:00:00 2001 From: amy-why-3459 Date: Fri, 17 Apr 2026 17:15:35 +0800 Subject: [PATCH 39/97] [BugFix] Fixing occasional engine crashes caused by abort requests (#2871) Signed-off-by: amy-why-3459 --- .../test_chunk_transfer_adapter.py | 51 +++++++++++++++++++ vllm_omni/core/sched/omni_ar_scheduler.py | 18 +++++++ .../core/sched/omni_generation_scheduler.py | 18 +++++++ .../chunk_transfer_adapter.py | 23 +++++++++ 4 files changed, 110 insertions(+) diff --git a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py index 7a3caba11..256e3e0a3 100644 --- a/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py +++ b/tests/distributed/omni_connectors/test_chunk_transfer_adapter.py @@ -4,10 +4,12 @@ import threading from collections import deque from types import SimpleNamespace +from unittest.mock import patch import pytest import torch from pytest_mock import MockerFixture +from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.request import RequestStatus from vllm_omni.distributed.omni_connectors.transfer_adapter.base import OmniTransferAdapterBase @@ -335,6 +337,27 @@ def test_cleanup_after_poll_flow(build_adapter): assert "ext-flow" not in adapter.request_payload +def test_finish_requests_restores_status(build_adapter): + """Abort path must pop ``requests_origin_status`` and restore pre-wait status. + + While ``process_pending_chunks`` holds a request off the scheduler queues, the + adapter records the prior status (WAITING or RUNNING). ``finish_requests`` must + put that status back on the live ``Request`` so base ``Scheduler.finish_requests`` + can finish bookkeeping without inconsistent state / crashes. + """ + adapter, _ = build_adapter(stage_id=1) + req_id = "req-abort-during-chunk" + prior = RequestStatus.RUNNING + request = _req(req_id, RequestStatus.WAITING_FOR_CHUNK) + adapter.requests_origin_status[req_id] = prior + requests_map = {req_id: request} + + adapter.finish_requests([req_id], RequestStatus.FINISHED_ABORTED, requests_map) + + assert request.status == prior + assert req_id not in adapter.requests_origin_status + + # --------------------------------------------------------------- # Scheduler trigger tests # --------------------------------------------------------------- @@ -508,3 +531,31 @@ def test_ar_scheduler_defers_cleanup_and_queues_save_on_finished(mocker: MockerF assert len(cleanup_calls) == 0 assert len(save_calls) == 1 + + +def test_omni_ar_scheduler_finish_requests(mocker: MockerFixture): + """``OmniARScheduler.finish_requests`` must run chunk adapter hook before vLLM base.""" + from vllm_omni.core.sched.omni_ar_scheduler import OmniARScheduler + + order: list[str] = [] + + adapter = mocker.MagicMock() + + def _adapter_finish(request_ids, finished_status, requests): + order.append("adapter") + return [] + + adapter.finish_requests.side_effect = _adapter_finish + + def _super_finish(_self, request_ids, finished_status): + order.append("super") + return [] + + sched = OmniARScheduler.__new__(OmniARScheduler) + sched.chunk_transfer_adapter = adapter + sched.requests = {} + + with patch.object(VLLMScheduler, "finish_requests", _super_finish): + OmniARScheduler.finish_requests(sched, ["r1"], RequestStatus.FINISHED_ABORTED) + + assert order == ["adapter", "super"] diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 0ee8cd16a..c5031ad11 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -539,6 +539,24 @@ class OmniARScheduler(VLLMScheduler): return engine_core_outputs + def finish_requests(self, request_ids: Any, finished_status: RequestStatus) -> list[tuple[str, int]]: + """Handles the finish signal from outside the scheduler. + + For example, the API server can abort a request when the client + disconnects. + + If request_ids is None, all requests will be finished. + + Returns: + Tuple of (req_id, client_index) for requests that were aborted. Will not + include any that were already finished. + """ + + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + + return super().finish_requests(request_ids, finished_status) + def _free_request(self, request: Request, delay_free_blocks: bool = False) -> dict[str, Any] | None: # TODO(wzliu)! for offline mode, we should not end process until all data is transferred """Mark a request as finished and free its resources.""" diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 1c4356d4f..87f2d282a 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -324,6 +324,24 @@ class OmniGenerationScheduler(VLLMScheduler): return scheduler_output + def finish_requests(self, request_ids, finished_status: RequestStatus) -> list[tuple[str, int]]: + """Handles the finish signal from outside the scheduler. + + For example, the API server can abort a request when the client + disconnects. + + If request_ids is None, all requests will be finished. + + Returns: + Tuple of (req_id, client_index) for requests that were aborted. Will not + include any that were already finished. + """ + + if self.chunk_transfer_adapter: + self.chunk_transfer_adapter.finish_requests(request_ids, finished_status, self.requests) + + return super().finish_requests(request_ids, finished_status) + """ Scheduler for the diffusion model. This scheduler is modified to stop the request immediately for the diffusion model. diff --git a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py index 393d0e801..1d5fc398a 100644 --- a/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py +++ b/vllm_omni/distributed/omni_connectors/transfer_adapter/chunk_transfer_adapter.py @@ -58,6 +58,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.waiting_for_chunk_waiting_requests: deque[Any] = deque() self.waiting_for_chunk_running_requests: deque[Any] = deque() self.requests_with_ready_chunks = set() + self.requests_origin_status = {} @classmethod def create_connector(cls, model_config: Any): @@ -279,6 +280,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.get_req_chunk.pop(request_id, None) self.requests_with_ready_chunks.discard(request_id) self.request_ids_mapping.pop(request_id, None) + self.requests_origin_status.pop(request_id, None) self._cancelled_load_reqs.add(request_id) self._finished_load_reqs.discard(request_id) @@ -408,6 +410,7 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): self.requests_with_ready_chunks.add(request.request_id) continue queue.remove(request) + self.requests_origin_status[request.request_id] = target_status waiting_for_chunk_list.append(request) def _clear_chunk_ready(self, scheduler_output: Any) -> None: @@ -420,3 +423,23 @@ class OmniChunkTransferAdapter(OmniTransferAdapterBase): for req_id in scheduler_output.scheduled_cached_reqs.req_ids: if req_id in self.requests_with_ready_chunks: self.requests_with_ready_chunks.remove(req_id) + + def finish_requests( + self, request_ids: Any, finished_status: RequestStatus, requests: dict[str, Request] | None = None + ) -> list[tuple[str, int]]: + assert RequestStatus.is_finished(finished_status) + if isinstance(request_ids, str): + request_ids = (request_ids,) + elif request_ids is not None: + request_ids = set(request_ids) + else: + request_ids = requests.keys() + + # First pass: collect requests to remove from queues + for req_id in request_ids: + request = requests.get(req_id) if requests else None + if request is None or request.is_finished(): + # Invalid request ID. + continue + if req_id in self.requests_origin_status: + request.status = self.requests_origin_status.pop(req_id) From a5a4998ece6374a76469ec3ac22e7b55c95e4e84 Mon Sep 17 00:00:00 2001 From: "Bingyu (Spencer) Liu" Date: Fri, 17 Apr 2026 18:43:17 +0800 Subject: [PATCH 40/97] [Feature] Support Prefill-Decode disaggregation via vLLM KV transfer (#2220) Signed-off-by: LiuBingyu --- tests/e2e/online_serving/test_qwen3_omni.py | 36 +- tests/entrypoints/test_pd_disaggregation.py | 1222 +++++++++++++++++ vllm_omni/engine/async_omni_engine.py | 46 + vllm_omni/engine/orchestrator.py | 132 +- vllm_omni/engine/output_processor.py | 5 +- vllm_omni/entrypoints/async_omni.py | 19 +- vllm_omni/entrypoints/omni.py | 16 +- vllm_omni/entrypoints/omni_base.py | 12 +- vllm_omni/entrypoints/pd_utils.py | 57 +- .../models/qwen3_omni/qwen3_omni.py | 32 +- .../stage_input_processors/qwen3_omni.py | 143 +- 11 files changed, 1672 insertions(+), 48 deletions(-) create mode 100644 tests/entrypoints/test_pd_disaggregation.py diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index 13af2ad11..9737fa42b 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -26,9 +26,15 @@ models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] QWEN3_OMNI_CONFIG_PATH = str(Path(__file__).parent.parent / "stage_configs" / "qwen3_omni_ci.yaml") QWEN3_OMNI_XPU_CONFIG_PATH = str(Path(__file__).parent.parent / "stage_configs" / "xpu" / "qwen3_omni_ci.yaml") +_STAGE_CONFIGS_DIR = Path(__file__).parent.parent / "stage_configs" +_PD_SEP_CONFIG = str(_STAGE_CONFIGS_DIR / "qwen3_omni_moe_pd_ci.yaml") -def get_chunk_config(config_path: str): - path = modify_stage_config( + +def get_chunk_config(config_path: str | None = None): + """Load qwen3_omni_ci.yaml with async_chunk modifications for streaming mode.""" + if config_path is None: + config_path = str(_STAGE_CONFIGS_DIR / "qwen3_omni_ci.yaml") + return modify_stage_config( config_path, updates={ "async_chunk": True, @@ -43,7 +49,6 @@ def get_chunk_config(config_path: str): }, deletes={"stage_args": {2: ["custom_process_input_func"]}}, ) - return path def get_prefix_caching_config(config_path: str): @@ -59,11 +64,20 @@ def get_prefix_caching_config(config_path: str): return path -if current_omni_platform.is_xpu(): - stage_configs = [QWEN3_OMNI_XPU_CONFIG_PATH] - prefix_caching_stage_configs = [get_prefix_caching_config(QWEN3_OMNI_XPU_CONFIG_PATH)] -else: # MI325 GPU should share the same config as H100 - stage_configs = [get_chunk_config(QWEN3_OMNI_CONFIG_PATH)] +# Set VLLM_TEST_PD_MODE=1 to test PD disaggregation, default tests async_chunk mode. +_USE_PD = os.environ.get("VLLM_TEST_PD_MODE", "0") == "1" + +# Stage configs for H100/CUDA, ROCm MI325, and XPU platforms +if current_omni_platform.is_rocm(): + rocm_config = str(_STAGE_CONFIGS_DIR / "rocm" / "qwen3_omni_ci.yaml") + stage_configs = [rocm_config] + prefix_caching_stage_configs = [get_prefix_caching_config(rocm_config)] +elif current_omni_platform.is_xpu(): + xpu_config = str(_STAGE_CONFIGS_DIR / "xpu" / "qwen3_omni_ci.yaml") + stage_configs = [xpu_config] + prefix_caching_stage_configs = [get_prefix_caching_config(xpu_config)] +else: + stage_configs = [_PD_SEP_CONFIG if _USE_PD else get_chunk_config(QWEN3_OMNI_CONFIG_PATH)] prefix_caching_stage_configs = [get_prefix_caching_config(QWEN3_OMNI_CONFIG_PATH)] # Create parameter combinations for model and stage config @@ -116,7 +130,8 @@ def get_max_batch_size(size_type="few"): @pytest.mark.advanced_model @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.") +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=3 if _USE_PD else 2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_mix_to_text_audio_001(omni_server, openai_client) -> None: """ @@ -155,7 +170,8 @@ def test_mix_to_text_audio_001(omni_server, openai_client) -> None: @pytest.mark.advanced_model @pytest.mark.core_model @pytest.mark.omni -@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=2) +@pytest.mark.skipif(_USE_PD, reason="Temporarily skip PD mode in this test module.") +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}, num_cards=3 if _USE_PD else 2) @pytest.mark.parametrize("omni_server", test_params, indirect=True) def test_text_to_text_001(omni_server, openai_client) -> None: """ diff --git a/tests/entrypoints/test_pd_disaggregation.py b/tests/entrypoints/test_pd_disaggregation.py new file mode 100644 index 000000000..5ffabfbf2 --- /dev/null +++ b/tests/entrypoints/test_pd_disaggregation.py @@ -0,0 +1,1222 @@ +"""Unit tests for PD (Prefill-Decode) disaggregation in the Omni orchestrator. + +Tests the PD detection, validation, config parsing, sampling param +preparation, and routing logic added by the PD disaggregation feature +(issue #1188). All tests run without GPU. + +NOTE (v1908 adaptation): Tests that relied on the old OmniStage / stage_list +architecture (removed in PR #1908) are marked xfail with +``reason="Requires migration to v1908 Orchestrator architecture"``. +The remaining tests exercise PDDisaggregationMixin directly and work +without spinning up a real engine. +""" + +import uuid +import warnings +from queue import Empty, Queue +from types import SimpleNamespace +from typing import Any + +import pytest +from vllm import SamplingParams + +from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin + +pytestmark = pytest.mark.skip(reason="Temporarily skip PD entrypoint tests while PD config is being removed.") + +# Suppress noisy DeprecationWarnings from optional Swig bindings imported by vLLM dependencies. +warnings.filterwarnings( + "ignore", + message=r"builtin type SwigPy.*has no __module__ attribute", + category=DeprecationWarning, +) + + +def _ns(**kwargs): + """Create a lightweight attribute object for tests.""" + return SimpleNamespace(**kwargs) + + +# --------------------------------------------------------------------------- +# Fake helpers (same pattern as test_omni_llm.py) +# --------------------------------------------------------------------------- + + +class _FakeEngineArgs(dict): + """Fake engine args that supports both attribute and dict access.""" + + def __init__(self, args_dict: dict[str, Any]): + super().__init__(args_dict) + if "model_stage" not in self: + self["model_stage"] = None + if "engine_output_type" not in self: + self["engine_output_type"] = None + for key, value in self.items(): + setattr(self, key, value) + + +class _FakeStageConfig: + def __init__(self, config_dict: dict[str, Any]): + engine_args_dict = config_dict.get("engine_args", {}) + self.engine_args = _FakeEngineArgs(engine_args_dict) + self.final_output = config_dict.get("final_output", False) + self.final_output_type = config_dict.get("final_output_type", None) + self.stage_id = config_dict.get("stage_id", 0) + self.is_prefill_only = config_dict.get("is_prefill_only", False) + self.is_decode_only = config_dict.get("is_decode_only", False) + self.engine_input_source = config_dict.get("engine_input_source", []) + self.is_comprehension = config_dict.get("is_comprehension", False) + self._config_dict = config_dict + + +class _FakeQueue: + def __init__(self, maxsize=0): + self._queue = Queue(maxsize=maxsize) + + def put(self, item): + self._queue.put(item) + + def put_nowait(self, item): + self._queue.put_nowait(item) + + def get(self): + return self._queue.get() + + def get_nowait(self): + return self._queue.get_nowait() + + def empty(self): + return self._queue.empty() + + +class _FakeStage: + """Lightweight stage stub with PD disaggregation flag support.""" + + def __init__(self, config, stage_init_timeout: int = 300): + if isinstance(config, dict): + config = _FakeStageConfig(config) + self.config = config + self.stage_config = config + self.engine = None + self.engine_outputs = None + self.stage_id = getattr(config, "stage_id", 0) + self.engine_args = config.engine_args + self.model_stage = getattr(config.engine_args, "model_stage", None) + self.stage_type = "llm" + self.default_sampling_params = SamplingParams(temperature=1.0) + self.final_output = config.final_output if hasattr(config, "final_output") else False + self.final_output_type = getattr(config, "final_output_type", None) + self.is_prefill_only = getattr(config, "is_prefill_only", False) + self.is_decode_only = getattr(config, "is_decode_only", False) + self.engine_input_source = getattr(config, "engine_input_source", []) + self.is_comprehension = getattr(config, "is_comprehension", False) + processed_input = getattr(config, "_config_dict", {}).get("processed_input", ["processed"]) + self._processed_input = processed_input + self._in_q = None + self._out_q = None + self._proc = None + self._stage_init_timeout = max(0, int(stage_init_timeout)) + + def attach_queues(self, in_q, out_q): + self._in_q = in_q + self._out_q = out_q + + def init_stage_worker( + self, model: str, *, is_async=False, shm_threshold_bytes=65536, ctx=None, batch_timeout=10, **kwargs + ): + self._proc = _ns( + start=lambda: None, + join=lambda timeout=None: None, + is_alive=lambda: False, + terminate=lambda: None, + ) + if self._out_q is not None: + try: + self._out_q.put_nowait({"type": "stage_ready", "stage_id": self.stage_id}) + except Exception: + pass + + def stop_stage_worker(self): + if self._in_q is not None: + try: + self._in_q.put_nowait({"type": "shutdown"}) + except Exception: + pass + + def submit(self, payload: dict[str, Any]): + if self._in_q is not None: + self._in_q.put(payload) + + def try_collect(self) -> Any: + if self._out_q is None: + return None + try: + return self._out_q.get_nowait() + except Empty: + return None + + def set_engine_outputs(self, outputs): + self.engine_outputs = outputs + + def process_engine_inputs(self, stage_list, prompts): + return self._processed_input + + +# --------------------------------------------------------------------------- +# Shared mock setup helpers +# --------------------------------------------------------------------------- + + +def _setup_engine_mocks(monkeypatch): + fake_engine = _ns() + fake_engine.tokenizer = _ns() + fake_engine.log_stats = False + fake_engine.vllm_config = _ns() + fake_engine.vllm_config.model_config = _ns() + fake_engine.vllm_config.model_config.io_processor_plugin = None + fake_engine.get_supported_tasks = lambda: [] + fake_engine.model_config = _ns() + fake_engine.model_config.io_processor_plugin = None + fake_registry = _ns() + fake_registry.resolve_model_cls = lambda *args, **kwargs: (_ns(), "test_arch") + fake_engine.model_config.registry = fake_registry + fake_engine.vllm_config.model_config.registry = fake_registry + + monkeypatch.setattr( + "vllm.v1.engine.llm_engine.LLMEngine.from_engine_args", + lambda **kw: fake_engine, + raising=False, + ) + + class FakeModelClass: + pass + + monkeypatch.setattr( + "vllm.model_executor.model_loader.utils.get_model_architecture", + lambda model_config: (FakeModelClass, "test_arch"), + raising=False, + ) + monkeypatch.setattr( + "vllm.model_executor.model_loader.utils._get_model_architecture", + lambda model_config: (FakeModelClass, "test_arch"), + raising=False, + ) + monkeypatch.setattr( + "vllm.model_executor.models.adapters.try_create_mm_pooling_model_cls", + lambda model_cls: model_cls, + raising=False, + ) + monkeypatch.setattr( + "vllm.multimodal.cache._enable_processor_cache", + lambda model_config, mm_registry: False, + raising=False, + ) + monkeypatch.setattr( + "vllm.plugins.io_processors.get_io_processor", + lambda vllm_config, io_processor_plugin: None, + raising=False, + ) + + +def _setup_multiprocessing_mocks(monkeypatch): + import multiprocessing as mp + + fake_process_instance = _ns( + start=lambda: None, + join=lambda timeout=None: None, + is_alive=lambda: False, + terminate=lambda: None, + ) + + def fake_process_class(*args, **kwargs): + return fake_process_instance + + fake_ctx = _ns() + fake_ctx.Queue = lambda maxsize=0: _FakeQueue(maxsize=maxsize) + fake_ctx.Process = fake_process_class + + monkeypatch.setattr(mp, "get_context", lambda method: fake_ctx, raising=False) + monkeypatch.setattr(mp, "Process", fake_process_class, raising=False) + + +def _setup_ipc_mocks(monkeypatch): + # These IPC helpers existed in the old architecture; no-op in new arch. + pass + + +def _setup_log_mocks(monkeypatch): + class _FakeOrchestratorAggregator: + def __init__(self, num_stages, enable_stats, wall_start_ts, final_stage_id_for_e2e=None): + self.num_stages = num_stages + self.enable_stats = enable_stats + self.stage_first_ts = [None] * num_stages + self.stage_last_ts = [None] * num_stages + self.stage_total_tokens = [0] * num_stages + self.accumulated_gen_time_ms = {} + self.e2e_done = set() + self.e2e_count = 0 + self.e2e_total_ms = 0.0 + + def on_stage_metrics(self, stage_id, req_id, metrics, final_output_type=None): + pass + + def on_finalize_request(self, stage_id, req_id, start_ts): + self.e2e_done.add(req_id) + + def on_forward(self, from_stage, to_stage, req_id, size_bytes, tx_ms, use_shm): + pass + + def accumulate_diffusion_metrics(self, stage_type, req_id, engine_outputs): + pass + + def record_audio_generated_frames(self, output, stage_id, req_id): + pass + + def stage_postprocess_timer(self, stage_id, req_id): + from contextlib import contextmanager + + @contextmanager + def _noop(): + yield + + return _noop() + + def build_and_log_summary(self): + return "Fake summary" + + monkeypatch.setattr( + "vllm_omni.entrypoints.omni.OrchestratorAggregator", + _FakeOrchestratorAggregator, + raising=False, + ) + + +def _clear_modules(): + import sys + + for module_name in [ + "vllm_omni.entrypoints.utils", + "vllm_omni.entrypoints.omni", + ]: + if module_name in sys.modules: + del sys.modules[module_name] + + +@pytest.fixture(autouse=True) +def mock_get_config(monkeypatch): + """Auto-mock get_config and related model loading functions.""" + import sys + + fake_tokenizer = _ns() + fake_tokenizer.encode = lambda *args, **kwargs: [1, 2, 3] + fake_tokenizer.decode = lambda *args, **kwargs: "test" + + def _mock_init_tokenizer_from_configs(model_config=None, **kwargs): + return fake_tokenizer + + monkeypatch.setattr( + "vllm.transformers_utils.tokenizer.init_tokenizer_from_configs", + _mock_init_tokenizer_from_configs, + raising=False, + ) + tokenizer_module_path = "vllm.transformers_utils.tokenizer" + if tokenizer_module_path in sys.modules: + setattr(sys.modules[tokenizer_module_path], "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) + + def _mock_length_from_prompt_token_ids_or_embeds(prompt_token_ids=None, prompt_embeds=None): + if prompt_token_ids is not None: + if isinstance(prompt_token_ids, list): + return len(prompt_token_ids) + return 10 + + monkeypatch.setattr( + "vllm.utils.length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds, raising=False + ) + monkeypatch.setattr( + "vllm_omni.engine.input_processor.length_from_prompt_token_ids_or_embeds", + _mock_length_from_prompt_token_ids_or_embeds, + raising=False, + ) + + processor_module_path = "vllm_omni.engine.input_processor" + if processor_module_path in sys.modules: + setattr( + sys.modules[processor_module_path], + "length_from_prompt_token_ids_or_embeds", + _mock_length_from_prompt_token_ids_or_embeds, + ) + + monkeypatch.setattr( + "vllm_omni.entrypoints.async_omni.init_tokenizer_from_configs", _mock_init_tokenizer_from_configs, raising=False + ) + async_omni_path = "vllm_omni.entrypoints.async_omni" + if async_omni_path in sys.modules: + setattr(sys.modules[async_omni_path], "init_tokenizer_from_configs", _mock_init_tokenizer_from_configs) + + fake_hf_config = _ns() + fake_hf_config.model_type = "qwen2_5_omni" + + monkeypatch.setattr( + "vllm.transformers_utils.config.get_config", lambda model, **kwargs: fake_hf_config, raising=False + ) + monkeypatch.setattr("vllm_omni.entrypoints.utils.get_config", lambda model, **kwargs: fake_hf_config, raising=False) + + def _mock_cached_file(path_or_repo_id, *args, **kwargs): + import os + import tempfile + + fake_config_file = os.path.join(tempfile.gettempdir(), "fake_config.json") + if not os.path.exists(fake_config_file): + with open(fake_config_file, "w") as f: + f.write('{"model_type": "qwen2_5_omni"}') + return fake_config_file + + monkeypatch.setattr("transformers.utils.hub.cached_file", _mock_cached_file, raising=False) + monkeypatch.setattr( + "transformers.utils.hub.cached_files", + lambda path_or_repo_id, filenames, **kwargs: ( + [_mock_cached_file(path_or_repo_id, filenames[0])] if filenames else None + ), + raising=False, + ) + + +# --------------------------------------------------------------------------- +# Helper to build an Omni instance with PD stage configs +# --------------------------------------------------------------------------- + + +def _make_pd_omni(monkeypatch, stage_configs, *, extra_setup=None): + """Create a lightweight PDDisaggregationMixin instance for unit tests. + + Bypasses the full OmniBase / AsyncOmniEngine init chain so tests run + without GPU. Returns an object that has all PDDisaggregationMixin + methods and state (``_pd_separation_pair``, ``_pd_kv_params_by_req``, + etc.) initialised from *stage_configs*. + + Tests that need the full ``Omni.generate()`` loop (old stage_list / queue + infrastructure) are marked ``xfail`` and not covered here. + """ + configs = [_FakeStageConfig(c) for c in stage_configs] + + class _LightweightOmni(PDDisaggregationMixin): + """Minimal shim: exposes stage_configs so PDDisaggregationMixin works.""" + + def __init__(self): + self._name = "Omni" + self._stage_configs = configs + self._init_pd_state() + + @property + def stage_configs(self): + return self._stage_configs + + if extra_setup: + import vllm_omni.entrypoints.omni as omni_module + + extra_setup(monkeypatch, omni_module) + + return _LightweightOmni() + + +# --------------------------------------------------------------------------- +# Stage config templates +# --------------------------------------------------------------------------- + + +def _prefill_stage_cfg(stage_id=0, **overrides): + cfg = { + "stage_id": stage_id, + "engine_args": { + "model_stage": "thinker", + "kv_transfer_config": { + "kv_connector": "MooncakeConnector", + "kv_role": "kv_producer", + "kv_rank": 0, + "kv_parallel_size": 2, + "kv_connector_extra_config": {"mooncake_bootstrap_port": 25201}, + }, + }, + "is_prefill_only": True, + "final_output": False, + "is_comprehension": True, + } + cfg.update(overrides) + return cfg + + +def _decode_stage_cfg(stage_id=1, engine_input_source=None, **overrides): + cfg = { + "stage_id": stage_id, + "engine_args": { + "model_stage": "thinker", + "kv_transfer_config": { + "kv_connector": "MooncakeConnector", + "kv_role": "kv_consumer", + "kv_rank": 1, + "kv_parallel_size": 2, + "kv_connector_extra_config": {"mooncake_bootstrap_port": 25202}, + }, + }, + "is_decode_only": True, + "engine_input_source": engine_input_source if engine_input_source is not None else [0], + "final_output": True, + "final_output_type": "text", + "is_comprehension": True, + } + cfg.update(overrides) + return cfg + + +def _talker_stage_cfg(stage_id=2, engine_input_source=None, **overrides): + cfg = { + "stage_id": stage_id, + "engine_args": {"model_stage": "talker"}, + "engine_input_source": engine_input_source if engine_input_source is not None else [1], + "final_output": False, + } + cfg.update(overrides) + return cfg + + +def _code2wav_stage_cfg(stage_id=3, engine_input_source=None, **overrides): + cfg = { + "stage_id": stage_id, + "engine_args": {"model_stage": "code2wav"}, + "engine_input_source": engine_input_source if engine_input_source is not None else [2], + "final_output": True, + "final_output_type": "audio", + } + cfg.update(overrides) + return cfg + + +# =================================================================== +# Tests: PD pair detection +# =================================================================== + + +class TestDetectPDSeparation: + """Tests for Omni._detect_pd_separation().""" + + def test_detects_pd_pair(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(stage_id=0), + _decode_stage_cfg(stage_id=1, engine_input_source=[0]), + ], + ) + assert omni._pd_separation_pair == (0, 1) + + def test_no_pd_pair_without_flags(self, monkeypatch): + """Normal (non-PD) pipeline has no PD pair.""" + omni = _make_pd_omni( + monkeypatch, + [ + { + "stage_id": 0, + "engine_args": {"model_stage": "thinker"}, + "final_output": True, + "final_output_type": "text", + }, + { + "stage_id": 1, + "engine_args": {"model_stage": "talker"}, + "engine_input_source": [0], + "final_output": True, + "final_output_type": "audio", + }, + ], + ) + assert omni._pd_separation_pair is None + + def test_detects_pd_pair_in_4_stage_pipeline(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(stage_id=0), + _decode_stage_cfg(stage_id=1, engine_input_source=[0]), + _talker_stage_cfg(stage_id=2, engine_input_source=[1]), + _code2wav_stage_cfg(stage_id=3, engine_input_source=[2]), + ], + ) + assert omni._pd_separation_pair == (0, 1) + + def test_pd_pair_uses_stage_id_for_input_source(self, monkeypatch): + """engine_input_source references stage_id, not list index.""" + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(stage_id=10), + _decode_stage_cfg(stage_id=20, engine_input_source=[10]), + ], + ) + assert omni._pd_separation_pair == (0, 1) + + +# =================================================================== +# Tests: PD config validation +# =================================================================== + + +class TestValidatePDConfig: + """Tests for Omni._validate_pd_separation_config().""" + + def test_valid_config_passes(self, monkeypatch): + """Valid PD config should not raise.""" + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + # If we got here without error, validation passed + assert omni._pd_separation_pair == (0, 1) + + def test_mismatched_connector_raises(self, monkeypatch): + """Different kv_connector types should raise ValueError.""" + decode_cfg = _decode_stage_cfg(engine_input_source=[0]) + decode_cfg["engine_args"]["kv_transfer_config"]["kv_connector"] = "NixlConnector" + + with pytest.raises(ValueError, match="connector mismatch"): + _make_pd_omni(monkeypatch, [_prefill_stage_cfg(), decode_cfg]) + + def test_wrong_prefill_role_raises(self, monkeypatch): + """Prefill with kv_consumer role should raise.""" + prefill_cfg = _prefill_stage_cfg() + prefill_cfg["engine_args"]["kv_transfer_config"]["kv_role"] = "kv_consumer" + + with pytest.raises(ValueError, match="kv_role must be"): + _make_pd_omni(monkeypatch, [prefill_cfg, _decode_stage_cfg(engine_input_source=[0])]) + + def test_wrong_decode_role_raises(self, monkeypatch): + """Decode with kv_producer role should raise.""" + decode_cfg = _decode_stage_cfg(engine_input_source=[0]) + decode_cfg["engine_args"]["kv_transfer_config"]["kv_role"] = "kv_producer" + + with pytest.raises(ValueError, match="kv_role must be"): + _make_pd_omni(monkeypatch, [_prefill_stage_cfg(), decode_cfg]) + + def test_missing_kv_transfer_config_raises(self, monkeypatch): + """Missing kv_transfer_config should raise.""" + prefill_cfg = _prefill_stage_cfg() + del prefill_cfg["engine_args"]["kv_transfer_config"] + + with pytest.raises(ValueError, match="kv_transfer_config"): + _make_pd_omni(monkeypatch, [prefill_cfg, _decode_stage_cfg(engine_input_source=[0])]) + + def test_mismatched_buffer_device_raises(self, monkeypatch): + """Mismatched kv_buffer_device should raise.""" + prefill_cfg = _prefill_stage_cfg() + prefill_cfg["engine_args"]["kv_transfer_config"]["kv_buffer_device"] = "cuda" + decode_cfg = _decode_stage_cfg(engine_input_source=[0]) + decode_cfg["engine_args"]["kv_transfer_config"]["kv_buffer_device"] = "cpu" + + with pytest.raises(ValueError, match="kv_buffer_device mismatch"): + _make_pd_omni(monkeypatch, [prefill_cfg, decode_cfg]) + + +# =================================================================== +# Tests: Connector info extraction +# =================================================================== + + +class TestGetPDConnectorInfo: + """Tests for Omni._get_pd_connector_info().""" + + def test_extracts_bootstrap_addr_for_mooncake(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + info = omni._pd_connector_info + assert "prefill_bootstrap_addr" in info + assert info["prefill_bootstrap_addr"] == "127.0.0.1:25201" + + def test_none_for_non_pd_pipeline(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + {"stage_id": 0, "engine_args": {}, "final_output": True, "final_output_type": "text"}, + ], + ) + assert omni._pd_connector_info is None + + +# =================================================================== +# Tests: Prefill sampling params preparation +# =================================================================== + + +class TestPreparePrefillSamplingParams: + """Tests for Omni._prepare_prefill_sampling_params().""" + + def test_sets_max_tokens_to_1(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048) + result = omni._prepare_prefill_sampling_params("req-1", sp) + + assert result.max_tokens == 1 + assert result is not sp # should be cloned + + def test_injects_kv_transfer_params(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048) + result = omni._prepare_prefill_sampling_params("req-1", sp) + + kv_params = result.extra_args["kv_transfer_params"] + assert kv_params["do_remote_decode"] is True + assert kv_params["do_remote_prefill"] is False + assert kv_params["transfer_id"] == "xfer-req-1" + + def test_preserves_existing_extra_args(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048, extra_args={"custom_key": "value"}) + result = omni._prepare_prefill_sampling_params("req-1", sp) + + assert result.extra_args["custom_key"] == "value" + assert "kv_transfer_params" in result.extra_args + + def test_does_not_mutate_original(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048) + _ = omni._prepare_prefill_sampling_params("req-1", sp) + + assert sp.max_tokens == 2048 + assert sp.extra_args is None + + +# =================================================================== +# Tests: Sampling params auto-duplication for PD split +# =================================================================== + + +@pytest.mark.xfail(reason="Requires migration to v1908 Orchestrator architecture (no stage_list / OmniStage)") +class TestSamplingParamsAutoDuplication: + """When user provides N-1 sampling params (for logical stages), the + orchestrator should auto-duplicate the thinker params for the decode stage. + """ + + def test_auto_duplicates_for_4_stage_pipeline(self, monkeypatch): + """User provides 3 params for 4 physical stages -> auto-insert decode params.""" + test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000001") + + def _extra_setup(mp, omni_module): + mp.setattr(uuid, "uuid4", lambda: test_uuid) + mp.setattr(omni_module, "uuid", uuid) + + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(stage_id=0), + _decode_stage_cfg(stage_id=1, engine_input_source=[0]), + _talker_stage_cfg(stage_id=2, engine_input_source=[1]), + _code2wav_stage_cfg(stage_id=3, engine_input_source=[2]), + ], + extra_setup=_extra_setup, + ) + + assert omni._pd_separation_pair == (0, 1) + assert len(omni.stage_list) == 4 + + # Simulate outputs for all stages + expected_rid = f"0_{test_uuid}" + for i in range(4): + omni.stage_list[i]._out_q.put_nowait( + { + "request_id": expected_rid, + "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1, 2])])], + "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0}, + } + ) + + # Provide 3 params (one less than 4 stages) - should auto-duplicate + sp_thinker = SamplingParams(temperature=0.4, max_tokens=2048) + sp_talker = SamplingParams(temperature=0.9, max_tokens=4096) + sp_code2wav = SamplingParams(temperature=0.0, max_tokens=65536) + + # This should NOT raise ValueError about param count mismatch + outputs = omni.generate( + prompts=["hello"], + sampling_params_list=[sp_thinker, sp_talker, sp_code2wav], + ) + assert isinstance(outputs, list) + + +# =================================================================== +# Tests: KV transfer params normalization +# =================================================================== + + +class TestNormalizeKVTransferParams: + def test_dict_passthrough(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + d = {"transfer_id": "test", "do_remote_decode": True} + assert omni._normalize_kv_transfer_params(d) is d + + def test_none_returns_none(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + assert omni._normalize_kv_transfer_params(None) is None + + def test_dataclass_to_dict(self, monkeypatch): + from dataclasses import dataclass + + @dataclass + class FakeKVParams: + transfer_id: str = "test" + do_remote_decode: bool = True + + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + result = omni._normalize_kv_transfer_params(FakeKVParams()) + assert isinstance(result, dict) + assert result["transfer_id"] == "test" + + +# =================================================================== +# Tests: _kv_cfg_to_dict +# =================================================================== + + +class TestKvCfgToDict: + def test_dict_passthrough(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + d = {"kv_connector": "MooncakeConnector"} + assert omni._kv_cfg_to_dict(d) is d + + def test_none_returns_empty(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + assert omni._kv_cfg_to_dict(None) == {} + + def test_dataclass_converted(self, monkeypatch): + from dataclasses import dataclass + + @dataclass + class FakeCfg: + kv_connector: str = "TestConnector" + kv_role: str = "kv_producer" + + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + result = omni._kv_cfg_to_dict(FakeCfg()) + assert result["kv_connector"] == "TestConnector" + assert result["kv_role"] == "kv_producer" + + +# =================================================================== +# Tests: PD routing in scheduling loop +# =================================================================== + + +@pytest.mark.xfail(reason="Requires migration to v1908 Orchestrator architecture (no stage_list / OmniStage)") +class TestPDRouting: + """Test that the scheduling loop correctly routes requests from + prefill to decode stage with proper kv_transfer_params. + """ + + def test_prefill_stage_receives_max_tokens_1(self, monkeypatch): + """Stage 0 (prefill) should receive max_tokens=1.""" + test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000002") + + def _extra_setup(mp, omni_module): + mp.setattr(uuid, "uuid4", lambda: test_uuid) + mp.setattr(omni_module, "uuid", uuid) + + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(stage_id=0), + _decode_stage_cfg(stage_id=1, engine_input_source=[0]), + ], + extra_setup=_extra_setup, + ) + + expected_rid = f"0_{test_uuid}" + + # Put stage outputs in both queues + omni.stage_list[0]._out_q.put_nowait( + { + "request_id": expected_rid, + "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1])])], + "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0}, + } + ) + omni.stage_list[1]._out_q.put_nowait( + { + "request_id": expected_rid, + "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1, 2, 3])])], + "metrics": {"num_tokens_out": 3, "stage_gen_time_ms": 50.0}, + } + ) + + sp_list = [SamplingParams(max_tokens=2048), SamplingParams(max_tokens=2048)] + omni.generate(prompts=["hello"], sampling_params_list=sp_list) + + # Check what was submitted to stage 0's input queue + # (skip the stage_ready message first) + task = omni.stage_list[0]._in_q.get_nowait() + assert task["sampling_params"].max_tokens == 1 + kv_params = task["sampling_params"].extra_args["kv_transfer_params"] + assert kv_params["do_remote_decode"] is True + assert kv_params["do_remote_prefill"] is False + assert kv_params["transfer_id"] == f"xfer-{expected_rid}" + + def test_decode_stage_receives_original_prompt(self, monkeypatch): + """Decode stage should get the original prompt (not processed outputs).""" + test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000003") + + def _extra_setup(mp, omni_module): + mp.setattr(uuid, "uuid4", lambda: test_uuid) + mp.setattr(omni_module, "uuid", uuid) + + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(stage_id=0), + _decode_stage_cfg(stage_id=1, engine_input_source=[0]), + ], + extra_setup=_extra_setup, + ) + + expected_rid = f"0_{test_uuid}" + original_prompt = "test prompt for PD" + + omni.stage_list[0]._out_q.put_nowait( + { + "request_id": expected_rid, + "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1])])], + "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0}, + } + ) + omni.stage_list[1]._out_q.put_nowait( + { + "request_id": expected_rid, + "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1, 2, 3])])], + "metrics": {"num_tokens_out": 3, "stage_gen_time_ms": 50.0}, + } + ) + + sp_list = [SamplingParams(max_tokens=2048), SamplingParams(max_tokens=2048)] + omni.generate(prompts=[original_prompt], sampling_params_list=sp_list) + + # Check what was forwarded to stage 1 (decode) + # The connector sends tasks to stage 1's input queue + task = omni.stage_list[1]._in_q.get_nowait() + # The engine_inputs should contain the original prompt + engine_inputs = task.get("engine_inputs") + # For PD routing, the original prompt is wrapped in a list + if isinstance(engine_inputs, list): + assert original_prompt in engine_inputs + else: + assert engine_inputs == original_prompt + + def test_decode_kv_params_have_correct_flags(self, monkeypatch): + """Decode stage kv_transfer_params should have correct role flags.""" + test_uuid = uuid.UUID("00000000-0000-0000-0000-000000000004") + + def _extra_setup(mp, omni_module): + mp.setattr(uuid, "uuid4", lambda: test_uuid) + mp.setattr(omni_module, "uuid", uuid) + + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(stage_id=0), + _decode_stage_cfg(stage_id=1, engine_input_source=[0]), + ], + extra_setup=_extra_setup, + ) + + expected_rid = f"0_{test_uuid}" + + omni.stage_list[0]._out_q.put_nowait( + { + "request_id": expected_rid, + "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1])])], + "metrics": {"num_tokens_out": 1, "stage_gen_time_ms": 10.0}, + } + ) + omni.stage_list[1]._out_q.put_nowait( + { + "request_id": expected_rid, + "engine_outputs": [_ns(request_id=expected_rid, outputs=[_ns(token_ids=[1, 2, 3])])], + "metrics": {"num_tokens_out": 3, "stage_gen_time_ms": 50.0}, + } + ) + + sp_list = [SamplingParams(max_tokens=2048), SamplingParams(max_tokens=2048)] + omni.generate(prompts=["hello"], sampling_params_list=sp_list) + + # Check decode task's kv_transfer_params + task = omni.stage_list[1]._in_q.get_nowait() + kv_params = task["sampling_params"].extra_args["kv_transfer_params"] + assert kv_params["do_remote_prefill"] is True + assert kv_params["do_remote_decode"] is False + assert kv_params["transfer_id"] == f"xfer-{expected_rid}" + assert kv_params["remote_bootstrap_addr"] == "127.0.0.1:25201" + + +# =================================================================== +# Tests: KV params cleanup +# =================================================================== + + +class TestKVParamsCleanup: + def test_drop_cleans_up(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + omni._pd_kv_params_by_req["req-1"] = {"transfer_id": "xfer-1"} + omni._drop_pd_kv_params("req-1") + assert "req-1" not in omni._pd_kv_params_by_req + + def test_drop_nonexistent_is_noop(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + omni._drop_pd_kv_params("nonexistent") # should not raise + + def test_pop_returns_stored_params(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + stored = {"transfer_id": "xfer-1", "extra_field": "value"} + omni._pd_kv_params_by_req["req-1"] = stored + + result = omni._pop_pd_kv_params("req-1") + assert result == stored + assert "req-1" not in omni._pd_kv_params_by_req + + def test_pop_uses_fallback_when_no_stored(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + fallback = {"transfer_id": "xfer-fallback"} + result = omni._pop_pd_kv_params("req-1", fallback=fallback) + assert result == fallback + + +# =================================================================== +# Tests: Config YAML loads without error +# =================================================================== + + +class TestPDYAMLConfig: + def test_pd_yaml_loads(self): + """The PD separation YAML config should load without errors.""" + import os + + yaml_path = os.path.join( + os.path.dirname(__file__), + "../../vllm_omni/model_executor/stage_configs/qwen3_omni_moe_pd_separation.yaml", + ) + yaml_path = os.path.abspath(yaml_path) + if not os.path.exists(yaml_path): + pytest.skip("PD separation YAML not found") + + from omegaconf import OmegaConf + + cfg = OmegaConf.load(yaml_path) + stages = cfg.stage_args + assert len(stages) == 4 + + # Prefill stage + assert stages[0].is_prefill_only is True + assert stages[0].final_output is False + assert stages[0].is_comprehension is True + + # Decode stage + assert stages[1].is_decode_only is True + assert stages[1].final_output is True + assert stages[1].final_output_type == "text" + assert stages[1].is_comprehension is True + assert 0 in stages[1].engine_input_source + + # KV transfer configs + assert stages[0].engine_args.kv_transfer_config.kv_role == "kv_producer" + assert stages[1].engine_args.kv_transfer_config.kv_role == "kv_consumer" + assert stages[0].engine_args.kv_transfer_config.kv_connector == "MooncakeConnector" + assert stages[1].engine_args.kv_transfer_config.kv_connector == "MooncakeConnector" + + +class TestPrefillStopNeutralization: + """Tests that _prepare_prefill_sampling_params neutralizes stop + conditions to ensure finish_reason='length'. + """ + + def test_clears_stop_strings(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048, stop=["", "STOP"]) + result = omni._prepare_prefill_sampling_params("req-1", sp) + assert result.stop == [] + + def test_clears_stop_token_ids(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048, stop_token_ids=[151643, 151644]) + result = omni._prepare_prefill_sampling_params("req-1", sp) + assert result.stop_token_ids == [] + + def test_clears_include_stop_str_in_output(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048, include_stop_str_in_output=True) + result = omni._prepare_prefill_sampling_params("req-1", sp) + assert result.include_stop_str_in_output is False + + def test_original_sp_unchanged(self, monkeypatch): + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + sp = SamplingParams(max_tokens=2048, stop=[""], stop_token_ids=[151643]) + _ = omni._prepare_prefill_sampling_params("req-1", sp) + assert sp.stop == [""] + assert sp.stop_token_ids == [151643] + + +# =================================================================== +# Tests: Failure mode & memory leak prevention +# =================================================================== +# NOTE: Full generate()-level failure mode tests are removed for now. +# The _run_generation error handler (line 1344-1350 in omni.py) calls +# _drop_pd_kv_params but does not increment completed_requests, causing +# the while-loop to hang. These tests need to be revisited once the +# production error-handling path is fixed to properly terminate on +# stage errors. + + +# =================================================================== +# Tests: TP size validation +# =================================================================== + + +class TestTPSizeValidation: + """Tests that _validate_pd_separation_config checks tensor_parallel_size.""" + + def test_matching_tp_passes(self, monkeypatch): + """Same TP size should not raise.""" + prefill_cfg = _prefill_stage_cfg() + prefill_cfg["engine_args"]["tensor_parallel_size"] = 2 + decode_cfg = _decode_stage_cfg(engine_input_source=[0]) + decode_cfg["engine_args"]["tensor_parallel_size"] = 2 + omni = _make_pd_omni(monkeypatch, [prefill_cfg, decode_cfg]) + assert omni._pd_separation_pair == (0, 1) + + def test_mismatched_tp_raises(self, monkeypatch): + """Different TP sizes should raise ValueError.""" + prefill_cfg = _prefill_stage_cfg() + prefill_cfg["engine_args"]["tensor_parallel_size"] = 2 + decode_cfg = _decode_stage_cfg(engine_input_source=[0]) + decode_cfg["engine_args"]["tensor_parallel_size"] = 4 + with pytest.raises(ValueError, match="tensor_parallel_size"): + _make_pd_omni(monkeypatch, [prefill_cfg, decode_cfg]) + + def test_default_tp_no_error(self, monkeypatch): + """Stages without explicit TP (defaults to 1) should pass.""" + omni = _make_pd_omni( + monkeypatch, + [ + _prefill_stage_cfg(), + _decode_stage_cfg(engine_input_source=[0]), + ], + ) + assert omni._pd_separation_pair == (0, 1) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index 238bbdcdb..fe6527349 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -25,6 +25,7 @@ from typing import TYPE_CHECKING, Any import janus import torch from omegaconf import OmegaConf +from vllm import envs as vllm_envs from vllm.engine.arg_utils import EngineArgs from vllm.inputs import PromptType from vllm.logger import init_logger @@ -79,6 +80,7 @@ from vllm_omni.engine.stage_init_utils import ( setup_stage_devices, terminate_alive_proc, ) +from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import ( inject_omni_kv_config, load_and_resolve_stage_configs, @@ -902,6 +904,7 @@ class AsyncOmniEngine: self._initialize_janus_queues() self._initialize_stages(stage_init_timeout) + pd_config = self._detect_pd_config() orchestrator = Orchestrator( request_async_queue=self.request_queue.async_q, output_async_queue=self.output_queue.async_q, @@ -910,6 +913,7 @@ class AsyncOmniEngine: stage_clients=self.stage_clients, output_processors=self.output_processors, stage_vllm_configs=self.stage_vllm_configs, + pd_config=pd_config, ) if not startup_future.done(): startup_future.set_result(asyncio.get_running_loop()) @@ -1131,6 +1135,48 @@ class AsyncOmniEngine: cache_config = AsyncOmniEngine._get_default_cache_config(cache_backend) return cache_config + def _detect_pd_config(self) -> dict[str, Any] | None: + """Detect PD (Prefill-Decode) disaggregation config from stage_configs. + Returns a dict with 'pd_pair' and 'bootstrap_addr', or None. + """ + pd_pair = PDDisaggregationMixin.detect_pd_separation_from_stage_configs(self.stage_configs) + if pd_pair is None: + return None + prefill_idx, decode_idx = pd_pair + + # Extract bootstrap address from prefill stage engine_args + bootstrap_addr: str | None = None + try: + prefill_cfg = self.stage_configs[prefill_idx] + ea = getattr(prefill_cfg, "engine_args", None) + kv_cfg = getattr(ea, "kv_transfer_config", None) if ea is not None else None + if kv_cfg is not None: + port = vllm_envs.VLLM_MOONCAKE_BOOTSTRAP_PORT + kv_ip = getattr(kv_cfg, "kv_ip", None) or "127.0.0.1" + bootstrap_addr = f"http://{kv_ip}:{port}" + except Exception as exc: + logger.warning("[AsyncOmniEngine] Could not extract PD bootstrap address: %s", exc) + + logger.info( + "[AsyncOmniEngine] PD disaggregation detected: prefill=stage-%d, decode=stage-%d, bootstrap=%s", + prefill_idx, + decode_idx, + bootstrap_addr, + ) + prefill_engine_id: str | None = None + try: + prefill_client = self.stage_clients[prefill_idx] + kv_cfg = getattr(getattr(prefill_client, "vllm_config", None), "kv_transfer_config", None) + prefill_engine_id = getattr(kv_cfg, "engine_id", None) + except Exception as exc: + logger.warning("[AsyncOmniEngine] Could not extract prefill engine_id: %s", exc) + + return { + "pd_pair": (prefill_idx, decode_idx), + "bootstrap_addr": bootstrap_addr, + "prefill_engine_id": prefill_engine_id, + } + @staticmethod def _create_default_diffusion_stage_cfg(kwargs: dict[str, Any]) -> list: """Create a default single-stage diffusion config from kwargs.""" diff --git a/vllm_omni/engine/orchestrator.py b/vllm_omni/engine/orchestrator.py index 1de4282fe..8204d70e6 100644 --- a/vllm_omni/engine/orchestrator.py +++ b/vllm_omni/engine/orchestrator.py @@ -42,6 +42,7 @@ def build_engine_core_request_from_tokens( params: SamplingParams | PoolingParams, arrival_time: float | None = None, model_config: ModelConfig | None = None, + mm_features: list | None = None, ) -> OmniEngineCoreRequest: """Build an OmniEngineCoreRequest directly from an OmniTokensPrompt. @@ -76,7 +77,7 @@ def build_engine_core_request_from_tokens( return OmniEngineCoreRequest( request_id=request_id, prompt_token_ids=prompt_token_ids, - mm_features=None, + mm_features=mm_features, sampling_params=sampling_params, pooling_params=pooling_params, arrival_time=arrival_time, @@ -104,6 +105,8 @@ class OrchestratorRequestState: # Metrics: timestamp when request was submitted to each stage stage_submit_ts: dict[int, float] = field(default_factory=dict) + mm_processor_kwargs: dict | None = None + mm_features: list | None = None class Orchestrator: @@ -123,6 +126,7 @@ class Orchestrator: stage_vllm_configs: list[Any], *, async_chunk: bool = False, + pd_config: dict[str, Any] | None = None, ) -> None: self.request_async_queue = request_async_queue self.output_async_queue = output_async_queue @@ -135,6 +139,16 @@ class Orchestrator: self.output_processors: list[Any] = output_processors self.stage_vllm_configs: list[Any] = stage_vllm_configs + # PD disaggregation state + self._pd_pair: tuple[int, int] | None = None + self._pd_bootstrap_addr: str | None = None + self._pd_prefill_engine_id: str | None = None + self._pd_kv_params: dict[str, Any] = {} + if pd_config is not None: + self._pd_pair = pd_config.get("pd_pair") + self._pd_bootstrap_addr = pd_config.get("bootstrap_addr") + self._pd_prefill_engine_id = pd_config.get("prefill_engine_id") + # Per-request state self.request_states: dict[str, OrchestratorRequestState] = {} @@ -359,6 +373,23 @@ class Orchestrator: } ) + # PD disaggregation: extract KV transfer params from prefill stage output + if self._pd_pair is not None and finished and stage_id == self._pd_pair[0]: + kv_params = getattr(output, "kv_transfer_params", None) + if kv_params is not None: + self._pd_kv_params[req_id] = kv_params if isinstance(kv_params, dict) else dict(kv_params) + logger.debug( + "[Orchestrator][PD] stored kv_transfer_params for req=%s (keys=%s)", + req_id, + list(self._pd_kv_params[req_id].keys()), + ) + else: + logger.warning( + "[Orchestrator][PD] prefill stage output for req=%s has no kv_transfer_params; " + "KV transfer may fail. Ensure apply_mooncake_connector_patch() was called.", + req_id, + ) + if ( finished and stage_id < req_state.final_stage_id @@ -371,6 +402,8 @@ class Orchestrator: await self._forward_to_next_stage(req_id, stage_id, output, req_state) if finished and stage_id == req_state.final_stage_id: + # PD: clean up any lingering KV params for this request + self._pd_kv_params.pop(req_id, None) self._cfg_tracker.cleanup_parent(req_id) self.request_states.pop(req_id, None) @@ -418,6 +451,51 @@ class Orchestrator: else: await self._forward_to_next_stage(req_id, stage_id, raw_output, req_state) + def _build_pd_decode_params(self, req_id: str, sp: Any) -> Any: + """Build decode-side sampling params with KV transfer params for PD routing. + + Clones the sampling params and injects kv_transfer_params that tell the + decode engine where to pull the KV cache from (prefill engine's bootstrap addr). + """ + sp = sp.clone() + if sp.extra_args is None: + sp.extra_args = {} + + # Get KV params captured from the prefill output (must include remote_request_id). + kv_prefill_params = self._pd_kv_params.pop(req_id, None) + if not kv_prefill_params or "remote_request_id" not in kv_prefill_params: + raise RuntimeError( + f"[Orchestrator][PD] Missing prefill kv_transfer_params.remote_request_id for req={req_id}" + ) + + decode_kv_params: dict[str, Any] = { + "transfer_id": f"xfer-{req_id}", + } + + if self._pd_bootstrap_addr: + decode_kv_params["remote_bootstrap_addr"] = self._pd_bootstrap_addr + + if self._pd_prefill_engine_id: + decode_kv_params["remote_engine_id"] = self._pd_prefill_engine_id + + # Overlay params from prefill side (includes remote_request_id set by monkey patch). + decode_kv_params.update(kv_prefill_params) + + # Ensure these flags are set correctly after any overlay. + decode_kv_params["do_remote_prefill"] = True + decode_kv_params["do_remote_decode"] = False + if not decode_kv_params.get("transfer_id"): + decode_kv_params["transfer_id"] = f"xfer-{req_id}" + + sp.extra_args["kv_transfer_params"] = decode_kv_params + + logger.debug( + "[Orchestrator][PD] decode kv_transfer_params for req=%s: %s", + req_id, + decode_kv_params, + ) + return sp + def _build_stage_metrics( self, stage_id: int, @@ -540,6 +618,52 @@ class Orchestrator: req_state.stage_submit_ts[next_stage_id] = _time.time() return + # PD disaggregation: prefill → decode routing uses original prompt + KV transfer params + if self._pd_pair is not None and (stage_id, next_stage_id) == self._pd_pair: + # Save prefill stage outputs so thinker2talker can merge embeddings later + self.stage_clients[stage_id].set_engine_outputs([output]) + + params = self._build_pd_decode_params(req_id, params) + + # Use the original user prompt for the decode stage (not processed embeddings) + original_prompt = req_state.prompt + raw_decode_inputs = [original_prompt] if not isinstance(original_prompt, list) else original_prompt + + decode_inputs: list[dict[str, Any]] = [] + for decode_input in raw_decode_inputs: + if isinstance(decode_input, dict): + decode_inputs.append(decode_input) + continue + prompt_token_ids = getattr(decode_input, "prompt_token_ids", None) + if prompt_token_ids is None: + raise TypeError( + "[Orchestrator][PD] decode input must be dict or have prompt_token_ids, " + f"got {type(decode_input).__name__} for req={req_id}" + ) + decode_inputs.append({"prompt_token_ids": list(prompt_token_ids)}) + + for decode_input in decode_inputs: + request = build_engine_core_request_from_tokens( + request_id=req_id, + prompt=decode_input, + params=params, + model_config=self.stage_vllm_configs[next_stage_id].model_config, + mm_features=req_state.mm_features, # Pass mm_features for M-RoPE + ) + request.external_req_id = request.request_id + + self.output_processors[next_stage_id].add_request( + request=request, + prompt=None, + parent_req=None, + request_index=0, + queue=None, + ) + await next_client.add_request_async(request) + + req_state.stage_submit_ts[next_stage_id] = _time.time() + return + self.stage_clients[stage_id].set_engine_outputs([output]) # Process inputs for next stage @@ -558,11 +682,16 @@ class Orchestrator: # Build and submit requests for each input for next_input in next_inputs: + # Only AR thinker stages consume encoder mm_features; downstream + # (talker/code2wav/…) must not see them (avoids encoder-cache misses). + _ms = getattr(next_client, "model_stage", None) + _mm_features = req_state.mm_features if _ms == "thinker" else None request = build_engine_core_request_from_tokens( request_id=req_id, prompt=next_input, params=params, model_config=self.stage_vllm_configs[next_stage_id].model_config, + mm_features=_mm_features, ) # TODO: Here we directly use the req id to assign. @@ -644,6 +773,7 @@ class Orchestrator: prompt=original_prompt, sampling_params_list=sampling_params_list, final_stage_id=final_stage_id, + mm_features=getattr(prompt, "mm_features", None), # Save mm_features for PD ) req_state.stage_submit_ts[stage_id] = _time.time() self.request_states[request_id] = req_state diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index badd799fc..67b4dd165 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -233,10 +233,9 @@ class OmniRequestState(RequestState): # Reuse base text/logprobs logic, then annotate with pooling_result. base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason, routed_experts) try: + if not hasattr(base_output, "multimodal_output"): + setattr(base_output, "multimodal_output", {}) if self.mm_accumulated is not None: - # Attach accumulated multimodal dict on the completion output - if not hasattr(base_output, "multimodal_output"): - setattr(base_output, "multimodal_output", {}) mm_out = getattr(base_output, "multimodal_output") if isinstance(mm_out, dict): for k, v in self.mm_accumulated.items(): diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 129ef3c99..5823aa4ab 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -78,7 +78,6 @@ class AsyncOmni(EngineClient, OmniBase): self.final_output_task: asyncio.Task | None = None self.config_path = self.engine.config_path - self.stage_configs = self.engine.stage_configs self.tts_max_instructions_length = kwargs.get("tts_max_instructions_length", None) self.input_processor = self.engine.input_processor @@ -209,6 +208,13 @@ class AsyncOmni(EngineClient, OmniBase): # Start final output dispatcher on the first call to generate() self._final_output_handler() + # Expand sampling params for PD disaggregation (user may provide N-1 params) + if ( + sampling_params_list is not None + and isinstance(sampling_params_list, Sequence) + and not isinstance(sampling_params_list, (str, bytes)) + ): + sampling_params_list = self._maybe_expand_sampling_params(list(sampling_params_list)) sampling_params_list = self.resolve_sampling_params_list(sampling_params_list) # Track per-request metrics @@ -228,20 +234,27 @@ class AsyncOmni(EngineClient, OmniBase): req_state.metrics = metrics self.request_states[request_id] = req_state + # PD disaggregation: modify prefill-stage sampling params per request + req_sp_list = list(sampling_params_list) + pd_pair = self._get_pd_separation_pair() + if pd_pair is not None: + p_id = pd_pair[0] + req_sp_list[p_id] = self._prepare_prefill_sampling_params(request_id, req_sp_list[p_id]) + # Add request(s) to stage 0. For streaming inputs, submit # chunks incrementally through streaming_update. if isinstance(prompt, AsyncGenerator): input_stream_task = await self._add_streaming_input_request( request_id=request_id, input_stream=prompt, - sampling_params_list=sampling_params_list, + sampling_params_list=req_sp_list, final_stage_id=final_stage_id_for_e2e, ) else: await self.engine.add_request_async( request_id=request_id, prompt=prompt, - sampling_params_list=sampling_params_list, + sampling_params_list=req_sp_list, final_stage_id=final_stage_id_for_e2e, ) submit_ts = time.time() diff --git a/vllm_omni/entrypoints/omni.py b/vllm_omni/entrypoints/omni.py index a3bfe98ce..8ef7e2ee5 100644 --- a/vllm_omni/entrypoints/omni.py +++ b/vllm_omni/entrypoints/omni.py @@ -66,6 +66,13 @@ class Omni(OmniBase): py_generator: bool = False, use_tqdm: bool | Callable[..., tqdm] = True, ) -> Generator[OmniRequestOutput, None, None] | list[OmniRequestOutput]: + # Expand sampling params for PD disaggregation (user may provide N-1 params) + if ( + sampling_params_list is not None + and isinstance(sampling_params_list, Sequence) + and not isinstance(sampling_params_list, (str, bytes)) + ): + sampling_params_list = self._maybe_expand_sampling_params(list(sampling_params_list)) sampling_params_list = self.resolve_sampling_params_list(sampling_params_list) try: if py_generator: @@ -125,10 +132,17 @@ class Omni(OmniBase): req_state.metrics = metrics self.request_states[req_id] = req_state + # PD disaggregation: modify stage-0 (prefill) sampling params per request + req_sp_list = list(sampling_params_list) + pd_pair = self._get_pd_separation_pair() + if pd_pair is not None: + p_id = pd_pair[0] + req_sp_list[p_id] = self._prepare_prefill_sampling_params(req_id, req_sp_list[p_id]) + self.engine.add_request( request_id=req_id, prompt=prompt, - sampling_params_list=sampling_params_list, + sampling_params_list=req_sp_list, final_stage_id=final_stage_id, ) submit_ts = time.time() diff --git a/vllm_omni/entrypoints/omni_base.py b/vllm_omni/entrypoints/omni_base.py index 1a7ffc4a5..82c648924 100644 --- a/vllm_omni/entrypoints/omni_base.py +++ b/vllm_omni/entrypoints/omni_base.py @@ -14,6 +14,7 @@ from vllm.v1.engine.exceptions import EngineDeadError from vllm_omni.engine.async_omni_engine import AsyncOmniEngine from vllm_omni.entrypoints.client_request_state import ClientRequestState +from vllm_omni.entrypoints.pd_utils import PDDisaggregationMixin from vllm_omni.entrypoints.utils import get_final_stage_id_for_e2e from vllm_omni.metrics.stats import OrchestratorAggregator as OrchestratorMetrics from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific @@ -65,7 +66,7 @@ def omni_snapshot_download(model_id: str) -> str: OutputMessageHandleResult = tuple[Literal[True], None, None, None] | tuple[Literal[False], str, int, ClientRequestState] -class OmniBase: +class OmniBase(PDDisaggregationMixin): """Shared runtime foundation for AsyncOmni and Omni.""" def __init__( @@ -84,6 +85,7 @@ class OmniBase: if "log_requests" in kwargs: raise TypeError("`log_requests` has been removed in Omni/AsyncOmni. Use `log_stats`.") model = omni_snapshot_download(model) + self._name = self.__class__.__name__ self.model = model self.log_stats = log_stats self.async_chunk = async_chunk @@ -125,10 +127,18 @@ class OmniBase: model, ) + # PD disaggregation state (detects if a prefill/decode stage pair is configured) + self._init_pd_state() + @property def num_stages(self) -> int: return self.engine.num_stages + @property + def stage_configs(self) -> list: + """Expose engine stage configs for PD disaggregation detection and validation.""" + return self.engine.stage_configs + @property def is_running(self) -> bool: return self.engine.is_alive() diff --git a/vllm_omni/entrypoints/pd_utils.py b/vllm_omni/entrypoints/pd_utils.py index 0e3d65f55..413d5d6b4 100644 --- a/vllm_omni/entrypoints/pd_utils.py +++ b/vllm_omni/entrypoints/pd_utils.py @@ -23,9 +23,19 @@ logger = logging.getLogger(__name__) class PDDisaggregationMixin: """Mixin supplying PD disaggregation helpers to OmniBase.""" + def _get_pd_separation_pair(self) -> tuple[int, int] | None: + """PD prefill/decode indices when ``_init_pd_state`` ran; else ``None``. + + Partial test doubles may skip ``OmniBase.__init__``; treat missing state as + no PD disaggregation instead of raising ``AttributeError``. + """ + return getattr(self, "_pd_separation_pair", None) + def _init_pd_state(self) -> None: """Initialise PD disaggregation state.""" - self._pd_separation_pair: tuple[int, int] | None = self._detect_pd_separation() + self._pd_separation_pair: tuple[int, int] | None = self.detect_pd_separation_from_stage_configs( + self.stage_configs + ) self._pd_connector_info: dict[str, Any] | None = None self._pd_kv_params_by_req: dict[str, dict[str, Any]] = {} self._pd_kv_params_lock = threading.Lock() @@ -40,11 +50,19 @@ class PDDisaggregationMixin: d_id, ) - def _detect_pd_separation(self) -> tuple[int, int] | None: - """Scan stage_list for a prefill/decode pair. Returns (p_id, d_id) or None.""" + @staticmethod + def detect_pd_separation_from_stage_configs(stage_configs: list[Any]) -> tuple[int, int] | None: + """Scan stage configs for a prefill/decode pair. + + Returns: + (prefill_idx, decode_idx) if one pair exists, None if not found. + + Raises: + ValueError: if multiple candidate PD pairs are found. + """ prefill_by_id: dict[int, int] = {} decode_indices: list[int] = [] - for i, stage in enumerate(self.stage_list): + for i, stage in enumerate(stage_configs): if getattr(stage, "is_prefill_only", False): prefill_by_id[i] = i sid = getattr(stage, "stage_id", i) @@ -55,7 +73,7 @@ class PDDisaggregationMixin: pd_pairs: list[tuple[int, int]] = [] for j in decode_indices: - source_ids = getattr(self.stage_list[j], "engine_input_source", []) + source_ids = getattr(stage_configs[j], "engine_input_source", []) for src in source_ids: if src in prefill_by_id: pd_pairs.append((prefill_by_id[src], j)) @@ -107,10 +125,11 @@ class PDDisaggregationMixin: def _validate_pd_separation_config(self) -> None: """Validate PD stage configurations are consistent.""" - assert self._pd_separation_pair is not None - p_id, d_id = self._pd_separation_pair - p_stage = self.stage_list[p_id] - d_stage = self.stage_list[d_id] + pair = self._get_pd_separation_pair() + assert pair is not None + p_id, d_id = pair + p_stage = self.stage_configs[p_id] + d_stage = self.stage_configs[d_id] def _get_kv_cfg(stage: "OmniStage") -> dict[str, Any]: ea = stage.engine_args @@ -158,11 +177,12 @@ class PDDisaggregationMixin: def _get_pd_connector_info(self) -> dict[str, Any] | None: """Extract prefill engine KV connector info.""" - if self._pd_separation_pair is None: + pair = self._get_pd_separation_pair() + if pair is None: return None - p_id, _ = self._pd_separation_pair - p_stage = self.stage_list[p_id] + p_id, _ = pair + p_stage = self.stage_configs[p_id] ea = p_stage.engine_args kv_cfg = getattr(ea, "kv_transfer_config", None) @@ -241,18 +261,17 @@ class PDDisaggregationMixin: def _is_pd_routing(self, stage_id: int, next_stage_id: int) -> bool: """True when edge stage_id → next_stage_id is the prefill→decode boundary.""" - return self._pd_separation_pair is not None and self._pd_separation_pair == ( - stage_id, - next_stage_id, - ) + pair = self._get_pd_separation_pair() + return pair is not None and pair == (stage_id, next_stage_id) def _maybe_expand_sampling_params(self, sampling_params_list: list) -> list: """Auto-duplicate thinker SP for decode stage when user provides N-1 params.""" - if self._pd_separation_pair is None: + pair = self._get_pd_separation_pair() + if pair is None: return sampling_params_list - if len(sampling_params_list) != len(self.stage_list) - 1: + if len(sampling_params_list) != len(self.stage_configs) - 1: return sampling_params_list - p_id, d_id = self._pd_separation_pair + p_id, d_id = pair sp_list = list(sampling_params_list) sp_list.insert(d_id, sp_list[p_id]) return sp_list diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 7df694797..cd8d147ca 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -891,10 +891,11 @@ class Qwen3OmniMoeForConditionalGeneration( Returns: (input_ids, input_embeds) for talker """ + target_len = thinker_result_ids.shape[-1] im_start_indexes = torch.cat( ( torch.nonzero(input_ids[0] == self.config.im_start_token_id).squeeze(), - torch.tensor([thinker_result_ids.shape[-1]], device=input_ids.device, dtype=input_ids.dtype), + torch.tensor([target_len], device=input_ids.device, dtype=input_ids.dtype), ), dim=-1, ) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here. @@ -1029,8 +1030,35 @@ class Qwen3OmniMoeForConditionalGeneration( return last_talker_hidden, text_step, update_dict def _get_talker_user_parts(self, im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed): + clamped = min( + segment_end_index, + multimodal_mask.shape[0], + thinker_hidden.shape[0], + thinker_embed.shape[0], + ) + if clamped < segment_end_index: + logger.warning( + "_get_talker_user_parts: segment_end_index %d clamped to %d " + "(embed=%d, hidden=%d, mask=%d). " + "This usually means _merge_pd_embeddings failed to merge " + "prefill embeddings – check PD prefill_mm keys.", + segment_end_index, + clamped, + thinker_embed.shape[0], + thinker_hidden.shape[0], + multimodal_mask.shape[0], + ) + segment_end_index = clamped + seg_len = segment_end_index - im_start_index + if seg_len <= 0: + return torch.empty( + (0, self.config.talker_config.text_config.hidden_size), + device=thinker_hidden.device, + dtype=torch.bfloat16, + ) + user_talker_part = torch.empty( - (segment_end_index - im_start_index, self.config.talker_config.text_config.hidden_size), + (seg_len, self.config.talker_config.text_config.hidden_size), device=thinker_hidden.device, dtype=torch.bfloat16, ) diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py index f4828fdda..c502041fe 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -3,6 +3,7 @@ # Copyright 2025 The Qwen team. """Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition.""" +import logging from typing import Any import torch @@ -18,6 +19,12 @@ from vllm_omni.model_executor.stage_input_processors.tts_utils import ( extract_speaker_from_request, ) +logger = logging.getLogger(__name__) + +# Pooling output layer keys: "0" = word embedding, "24" = accept_hidden_layer +_EMBED_LAYER_KEY = "0" +_HIDDEN_LAYER_KEY = "24" + def _compute_talker_prompt_ids_length(info, device: torch.device | str = "cuda") -> int: im_start_token_id = 151644 @@ -84,6 +91,95 @@ def _validate_stage_inputs(stage_list, engine_input_source): return stage.engine_outputs +# ========================= +# PD disaggregation helpers +# ========================= + + +def _get_prefill_stage(stage_list: list[Any], source_stage_id: int) -> Any | None: + if source_stage_id <= 0: + return None + source_stage = stage_list[source_stage_id] + if not getattr(source_stage, "is_decode_only", False): + return None + prev_stage = stage_list[source_stage_id - 1] + if getattr(prev_stage, "is_prefill_only", False) and prev_stage.engine_outputs is not None: + return prev_stage + return None + + +def _merge_pd_embeddings( + decode_emb: torch.Tensor, + decode_hid: torch.Tensor, + prefill_mm: dict[str, Any], + device: torch.device, + expected_total: int | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + """Merge prefill prompt embeddings with decode generated embeddings. + + In PD mode the prefill engine processes the prompt and the decode engine + generates tokens starting from position 1. This function concatenates + them, removing the overlapping token(s): + + merged = prefill[:P] + decode[overlap:] + + where overlap = P + D - expected_total. + """ + try: + p_emb = prefill_mm[_EMBED_LAYER_KEY].detach().to(device=device, dtype=torch.float) + p_hid = prefill_mm[_HIDDEN_LAYER_KEY].detach().to(device=device, dtype=torch.float) + except (KeyError, AttributeError, TypeError) as exc: + available_keys = list(prefill_mm.keys()) if isinstance(prefill_mm, dict) else type(prefill_mm).__name__ + logger.error( + "_merge_pd_embeddings: failed to extract prefill embeddings (%s). " + "Expected keys %r and %r, got: %s. " + "Falling back to decode-only embeddings – talker user-segment will be degraded.", + exc, + _EMBED_LAYER_KEY, + _HIDDEN_LAYER_KEY, + available_keys, + ) + return decode_emb, decode_hid + + if p_emb.shape[0] == 0 or decode_emb.shape[0] == 0: + return decode_emb, decode_hid + + raw_total = p_emb.shape[0] + decode_emb.shape[0] + overlap = max(0, raw_total - expected_total) if expected_total is not None else 0 + + merged_emb = torch.cat([p_emb, decode_emb[overlap:]], dim=0) + merged_hid = torch.cat([p_hid, decode_hid[overlap:]], dim=0) + return merged_emb, merged_hid + + +def _get_prefill_multimodal_output(prefill_stage: Any, output_index: int) -> dict[str, Any] | None: + """Return multimodal_output dict from the PD prefill stage for a given batch index.""" + try: + prefill_eos = prefill_stage.engine_outputs + prefill_eo = prefill_eos[min(output_index, len(prefill_eos) - 1)] + return prefill_eo.outputs[0].multimodal_output + except Exception: + return None + + +def _resolve_tts_token_embedding( + key: str, + *, + thinker_mm: dict[str, Any], + prefill_mm: dict[str, Any] | None, + device: torch.device, +) -> torch.Tensor | None: + """Return TTS BOS/EOS/PAD embedding tensors for the talker projection path. + + Values are taken from the current thinker (decode) ``multimodal_output``; in + PD mode, missing keys may be filled from the paired prefill stage output. + """ + val = thinker_mm.get(key) + if val is None and prefill_mm is not None: + val = prefill_mm.get(key) + return val.detach().to(device=device, dtype=torch.float) if val is not None else None + + # ========================= # Thinker -> Talker # ========================= @@ -111,8 +207,8 @@ def thinker2talker_async_chunk( all_token_ids = _ensure_list(all_token_ids) prompt_token_ids = _ensure_list(prompt_token_ids) talker_additional_info = { - "thinker_prefill_embeddings": pooling_output.get("0").detach().cpu(), - "thinker_hidden_states": pooling_output.get("24").detach().cpu(), + "thinker_prefill_embeddings": pooling_output.get(_EMBED_LAYER_KEY).detach().cpu(), + "thinker_hidden_states": pooling_output.get(_HIDDEN_LAYER_KEY).detach().cpu(), "thinker_sequences": all_token_ids, "thinker_input_ids": prompt_token_ids, # Provide thinker-side TTS token embeddings for talker projection @@ -161,7 +257,7 @@ def thinker2talker_async_chunk( if output_token_ids: talker_additional_info["override_keys"] = ["thinker_decode_embeddings", "thinker_output_token_ids"] - talker_additional_info["thinker_decode_embeddings"] = pooling_output.get("0").detach().cpu() + talker_additional_info["thinker_decode_embeddings"] = pooling_output.get(_EMBED_LAYER_KEY).detach().cpu() talker_additional_info["thinker_output_token_ids"] = output_token_ids else: # When prefilling a chunked thinker, thinker_hidden_states needs to be updated. @@ -185,6 +281,9 @@ def thinker2talker( 2. Split hidden states into: prompt embeddings + generated embeddings 3. Package for talker with additional information + In PD disaggregation mode, merges prefill-stage prompt embeddings with + decode-stage generated embeddings before handing off to the talker. + Args: stage_list: List of stage objects engine_input_source: Source stage IDs (typically [0] for thinker) @@ -199,21 +298,49 @@ def thinker2talker( device = torch.device(current_platform.device_type) + # PD disaggregation: look up the preceding prefill stage (if any) + source_stage_id = engine_input_source[0] + prefill_stage = _get_prefill_stage(stage_list, source_stage_id) + # Process each thinker output for i, thinker_output in enumerate(thinker_outputs): output = thinker_output.outputs[0] + thinker_mm = output.multimodal_output + # Full thinker embedding sequence for the talker: single thinker engine in the + # non-PD path; after optional merge with prefill-side tensors in PD mode. + thinker_emb = thinker_mm[_EMBED_LAYER_KEY].detach().to(device=device, dtype=torch.float) + thinker_hid = thinker_mm[_HIDDEN_LAYER_KEY].detach().to(device=device, dtype=torch.float) + + prefill_mm: dict[str, Any] | None = None + if prefill_stage is not None: + prefill_mm = _get_prefill_multimodal_output(prefill_stage, i) + + if prefill_mm is not None: + expected_total = len(thinker_output.prompt_token_ids) + len(output.token_ids) + try: + thinker_emb, thinker_hid = _merge_pd_embeddings( + thinker_emb, thinker_hid, prefill_mm, device, expected_total=expected_total + ) + except Exception as exc: + logger.warning("[PD] Could not merge prefill embeddings: %s", exc) info = { - "thinker_prefill_embeddings": output.multimodal_output["0"].detach().to(device=device, dtype=torch.float), - "thinker_hidden_states": output.multimodal_output["24"].detach().to(device=device, dtype=torch.float), + "thinker_prefill_embeddings": thinker_emb, + "thinker_hidden_states": thinker_hid, "thinker_sequences": ( thinker_output.prompt_token_ids + output.token_ids ), # the thinker_sequences is the whole ids "thinker_input_ids": thinker_output.prompt_token_ids, # Provide thinker-side TTS token embeddings for talker projection - "tts_bos_embed": output.multimodal_output["tts_bos_embed"].detach().to(device=device, dtype=torch.float), - "tts_eos_embed": output.multimodal_output["tts_eos_embed"].detach().to(device=device, dtype=torch.float), - "tts_pad_embed": output.multimodal_output["tts_pad_embed"].detach().to(device=device, dtype=torch.float), + "tts_bos_embed": _resolve_tts_token_embedding( + "tts_bos_embed", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device + ), + "tts_eos_embed": _resolve_tts_token_embedding( + "tts_eos_embed", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device + ), + "tts_pad_embed": _resolve_tts_token_embedding( + "tts_pad_embed", thinker_mm=thinker_mm, prefill_mm=prefill_mm, device=device + ), } speaker = extract_speaker_from_prompt(prompt, index=i) if speaker is not None: From c0ccbb872f018851f2c5a6c168e8175eeb10704e Mon Sep 17 00:00:00 2001 From: Yuanheng Zhao <54058983+yuanheng-zhao@users.noreply.github.com> Date: Fri, 17 Apr 2026 19:10:02 +0800 Subject: [PATCH 41/97] [Model] Add Ming-flash-omni-2.0 Thinker Stage (#1822) Signed-off-by: yuanheng Signed-off-by: Yuanheng Zhao --- .../ming_flash_omni/README.md | 76 ++ .../ming_flash_omni/end2end.py | 485 ++++++++++ .../online_serving/ming_flash_omni/README.md | 204 ++++ .../run_curl_multimodal_generation.sh | 145 +++ tests/conftest.py | 4 + .../offline_inference/test_ming_flash_omni.py | 142 +++ .../online_serving/test_ming_flash_omni.py | 247 +++++ .../bailingmm_moe_v2_lite_ci.yaml | 35 + .../models/ming_flash_omni/__init__.py | 18 + .../models/ming_flash_omni/audio_encoder.py | 246 +++++ .../models/ming_flash_omni/ming_flash_omni.py | 223 +++++ .../ming_flash_omni_thinker.py | 893 +++++++++++++++++ .../modeling_bailing_moe_v2.py | 896 ++++++++++++++++++ .../models/ming_flash_omni/projectors.py | 184 ++++ .../models/ming_flash_omni/vision_encoder.py | 125 +++ .../qwen3_tts/tokenizer_25hz/vq/speech_vq.py | 3 +- .../tokenizer_25hz/vq/whisper_encoder.py | 25 +- vllm_omni/model_executor/models/registry.py | 17 + .../model_executor/models/whisper_utils.py | 39 + .../stage_configs/bailingmm_moe_v2_lite.yaml | 46 + .../transformers_utils/configs/__init__.py | 11 + .../configs/ming_flash_omni.py | 302 ++++++ .../transformers_utils/processors/__init__.py | 12 + .../transformers_utils/processors/ming.py | 430 +++++++++ 24 files changed, 4783 insertions(+), 25 deletions(-) create mode 100644 examples/offline_inference/ming_flash_omni/README.md create mode 100644 examples/offline_inference/ming_flash_omni/end2end.py create mode 100644 examples/online_serving/ming_flash_omni/README.md create mode 100755 examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh create mode 100644 tests/e2e/offline_inference/test_ming_flash_omni.py create mode 100644 tests/e2e/online_serving/test_ming_flash_omni.py create mode 100644 tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/__init__.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/audio_encoder.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/modeling_bailing_moe_v2.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/projectors.py create mode 100644 vllm_omni/model_executor/models/ming_flash_omni/vision_encoder.py create mode 100644 vllm_omni/model_executor/models/whisper_utils.py create mode 100644 vllm_omni/model_executor/stage_configs/bailingmm_moe_v2_lite.yaml create mode 100644 vllm_omni/transformers_utils/configs/ming_flash_omni.py create mode 100644 vllm_omni/transformers_utils/processors/__init__.py create mode 100644 vllm_omni/transformers_utils/processors/ming.py diff --git a/examples/offline_inference/ming_flash_omni/README.md b/examples/offline_inference/ming_flash_omni/README.md new file mode 100644 index 000000000..7414163fc --- /dev/null +++ b/examples/offline_inference/ming_flash_omni/README.md @@ -0,0 +1,76 @@ +# Ming-flash-omni 2.0 + +[Ming-flash-omni-2.0](https://github.com/inclusionAI/Ming) is an omni-modal model supporting text, image, video, and audio understanding, with outputs in text, image, and audio. For now, Ming-flash-omni-2.0 in vLLM-Omni is supported with thinker stage (multi-modal understanding). + +## Setup + +Please refer to the [stage configuration documentation](https://docs.vllm.ai/projects/vllm-omni/en/latest/configuration/stage_configs/) to configure memory allocation appropriately for your hardware setup. + +## Run examples + +### Text-only +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type text +``` + +#### Reasoning (Thinking Mode) + +Reasoning (Thinking) mode is enabled via applying "detailed thinking on" when building the system prompt template (in `apply_chat_template`). + +In the end2end example, a default problem for thinking mode is provided, as referred to the example usage of Ming's cookbook; +To utilize it, you have to download the example figure from https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/3_0.png + +```bash +python examples/offline_inference/ming_flash_omni/end2end.py -q reasoning --image-path ./3_0.png +``` + +### Image understanding +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image + +# With a local image +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image --image-path /path/to/image.jpg +``` + +### Audio understanding +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio + +# With a local audio file +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --audio-path /path/to/audio.wav +``` + +### Video understanding +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video + +# With a local video and custom frame count +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_video --video-path /path/to/video.mp4 --num-frames 16 +``` + +### Mixed modalities (image + audio) +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_mixed_modalities \ + --image-path /path/to/image.jpg \ + --audio-path /path/to/audio.wav +``` + +If media file paths are not provided, the script uses built-in default assets. + +### Modality control +To control output modalities (e.g. text-only output): +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_audio --modalities text +``` + +*For now, only text output is supported* + +### Custom stage config +```bash +python examples/offline_inference/ming_flash_omni/end2end.py --query-type use_image \ + --stage-configs-path /path/to/your_config.yaml +``` + +## Online serving + +For online serving via the OpenAI-compatible API, see [examples/online_serving/ming_flash_omni/README.md](../../online_serving/ming_flash_omni/README.md). diff --git a/examples/offline_inference/ming_flash_omni/end2end.py b/examples/offline_inference/ming_flash_omni/end2end.py new file mode 100644 index 000000000..49cdbcc01 --- /dev/null +++ b/examples/offline_inference/ming_flash_omni/end2end.py @@ -0,0 +1,485 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Partial example cases are referred from +# https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/cookbook.ipynb +import os +import time +from typing import NamedTuple + +import librosa +import numpy as np +import vllm +from PIL import Image +from transformers import AutoProcessor +from vllm import SamplingParams +from vllm.assets.audio import AudioAsset +from vllm.assets.image import ImageAsset +from vllm.assets.video import VideoAsset, video_to_ndarrays +from vllm.multimodal.image import convert_image_mode +from vllm.utils.argparse_utils import FlexibleArgumentParser + +import vllm_omni +from vllm_omni.entrypoints.omni import Omni + +# Imports the processor also registers itself +from vllm_omni.transformers_utils.processors.ming import MingFlashOmniProcessor # noqa: F401 + +SEED = 42 +MODEL_NAME = "Jonathan1909/Ming-flash-omni-2.0" + + +class QueryResult(NamedTuple): + inputs: dict + limit_mm_per_prompt: dict[str, int] + + +def get_text_query(processor: MingFlashOmniProcessor, question: str | None = None) -> QueryResult: + if question is None: + question = "请详细介绍鹦鹉的生活习性。" + conversation = [{"role": "HUMAN", "content": question}] + prompt = processor.apply_chat_template(conversation, tokenize=False) + return QueryResult( + inputs={"prompt": prompt}, + limit_mm_per_prompt={}, + ) + + +def get_image_query( + processor: MingFlashOmniProcessor, + question: str | None = None, + image_path: str | None = None, +) -> QueryResult: + if question is None: + question = "Describe this image in detail." + + if image_path: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + image_data = convert_image_mode(Image.open(image_path), "RGB") + else: + image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB") + + conversation = [ + { + "role": "HUMAN", + "content": [ + {"type": "image", "image": image_data}, + {"type": "text", "text": question}, + ], + } + ] + prompt = processor.apply_chat_template(conversation, tokenize=False) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": {"image": image_data}, + }, + limit_mm_per_prompt={"image": 1}, + ) + + +def get_audio_query( + processor: MingFlashOmniProcessor, + question: str | None = None, + audio_path: str | None = None, + sampling_rate: int = 16000, +) -> QueryResult: + if question is None: + question = "Please recognize the language of this speech and transcribe it. Format: oral." + + if audio_path: + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + audio_signal, sr = librosa.load(audio_path, sr=sampling_rate) + audio_data = (audio_signal.astype(np.float32), sr) + else: + audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate + + # Use a string for "audio" so the processor counts it as 1 audio input + conversation = [ + { + "role": "HUMAN", + "content": [ + {"type": "audio", "audio": "input"}, + {"type": "text", "text": question}, + ], + } + ] + prompt = processor.apply_chat_template(conversation, tokenize=False) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": {"audio": audio_data}, + }, + limit_mm_per_prompt={"audio": 1}, + ) + + +def get_video_query( + processor: MingFlashOmniProcessor, + question: str | None = None, + video_path: str | None = None, + num_frames: int = 16, +) -> QueryResult: + if question is None: + question = "Describe what is happening in this video." + + if video_path: + if not os.path.exists(video_path): + raise FileNotFoundError(f"Video file not found: {video_path}") + video_frames = video_to_ndarrays(video_path, num_frames=num_frames) + else: + video_frames = VideoAsset(name="baby_reading", num_frames=num_frames).np_ndarrays + + conversation = [ + { + "role": "HUMAN", + "content": [ + {"type": "video"}, + {"type": "text", "text": question}, + ], + } + ] + prompt = processor.apply_chat_template(conversation, tokenize=False) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": {"video": video_frames}, + }, + limit_mm_per_prompt={"video": 1}, + ) + + +def get_mixed_modalities_query( + processor: MingFlashOmniProcessor, + image_path: str | None = None, + audio_path: str | None = None, + sampling_rate: int = 16000, +) -> QueryResult: + """Mixed image + audio understanding.""" + question = "Describe the image, and recognize the language of this speech and transcribe it. Format: oral" + + if image_path: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + image_data = convert_image_mode(Image.open(image_path), "RGB") + else: + image_data = convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB") + + if audio_path: + if not os.path.exists(audio_path): + raise FileNotFoundError(f"Audio file not found: {audio_path}") + sig, sr = librosa.load(audio_path, sr=sampling_rate) + audio_data = (sig.astype(np.float32), sr) + else: + audio_data = AudioAsset("mary_had_lamb").audio_and_sample_rate + + conversation = [ + { + "role": "HUMAN", + "content": [ + {"type": "image", "image": image_data}, + {"type": "audio", "audio": "input"}, + {"type": "text", "text": question}, + ], + } + ] + prompt = processor.apply_chat_template(conversation, tokenize=False) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": {"image": image_data, "audio": audio_data}, + }, + limit_mm_per_prompt={"image": 1, "audio": 1}, + ) + + +def get_reasoning_query( + processor: MingFlashOmniProcessor, + question: str | None = None, + image_path: str | None = None, +) -> QueryResult: + if question is None: + # NOTE: To use the following default question, input with example figure provided by Ming + # https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/figures/cases/3_0.png + # E.g., + # python examples/offline_inference/ming_flash_omni/end2end.py -q reasoning --image-path ./3_0.png + # Otherwise, the problem solving might be false. + question = ( + "Based on the following rules:\n•\tYou control the smiley face character\n" + "•\tYou can move up, down, left, and right, and only a single square at a time\n" + "•\tWalls are dark grey and cannot be moved into\n•\tThe brown square is a box\n•" + "\tThe box can be pushed by moving into it (i.e., if you are in the square " + "adjacent to the box to the left, and move onto the square with the box, " + "the box will move one square to the right).\n" + "•\tThe box cannot be pushed into walls\n" + "•\tThe blue door at the bottom is locked and cannot be passed through, " + "unless the box is placed on the blue square\n" + "•\tThe square beneath the blue door is the exit\n" + "•\tMoving from one square to another\n\n" + "Let's assume a coordinate system where the smiley face is " + "on the top left at (1,1) and the square below it is (1,2). " + "The smiley face performs the following moves: {down, right, right, right}, " + "such that the smiley face is at square (4,2) and the box is in square (5,2). " + "What are the next sequence of moves that must be done to move the box down to (5,3)? " + "Give your answer as a comma separated list." + ) + + if image_path: + if not os.path.exists(image_path): + raise FileNotFoundError(f"Image file not found: {image_path}") + image_data = convert_image_mode(Image.open(image_path), "RGB") + conversation = [ + { + "role": "HUMAN", + "content": [ + {"type": "image", "image": image_data}, + {"type": "text", "text": question}, + ], + } + ] + prompt = processor.apply_chat_template(conversation, tokenize=False, use_cot_system_prompt=True) + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": {"image": image_data}, + }, + limit_mm_per_prompt={"image": 1}, + ) + + conversation = [{"role": "HUMAN", "content": question}] + prompt = processor.apply_chat_template(conversation, tokenize=False, use_cot_system_prompt=True) + return QueryResult( + inputs={"prompt": prompt}, + limit_mm_per_prompt={}, + ) + + +query_map = { + "text": get_text_query, + "use_audio": get_audio_query, + "use_image": get_image_query, + "use_video": get_video_query, + "use_mixed_modalities": get_mixed_modalities_query, + "reasoning": get_reasoning_query, +} + + +def main(args): + print( + "=" * 20, + "\n", + f"vllm version: {vllm.__version__}\n", + f"vllm-omni version: {vllm_omni.__version__}\n", + "=" * 20, + sep="", + ) + + processor = AutoProcessor.from_pretrained(MODEL_NAME, trust_remote_code=True) + assert isinstance(processor, MingFlashOmniProcessor), f"Wrong processor type being used: {type(processor)}" + + query_func = query_map[args.query_type] + if args.query_type == "use_image": + query_result = query_func(processor, image_path=args.image_path) + elif args.query_type == "use_audio": + query_result = query_func(processor, audio_path=args.audio_path, sampling_rate=args.sampling_rate) + elif args.query_type == "use_video": + query_result = query_func(processor, video_path=args.video_path, num_frames=args.num_frames) + elif args.query_type == "use_mixed_modalities": + query_result = query_func( + processor, + image_path=args.image_path, + audio_path=args.audio_path, + sampling_rate=args.sampling_rate, + ) + elif args.query_type == "reasoning": + query_result = query_func(processor, image_path=args.image_path) + else: + query_result = query_func(processor) + + # Initialize Omni (with thinker-only stage config) + omni = Omni( + model=MODEL_NAME, + stage_configs_path=args.stage_configs_path, + log_stats=args.log_stats, + init_timeout=args.init_timeout, + stage_init_timeout=args.stage_init_timeout, + ) + + # Thinker sampling params + thinker_sampling_params = SamplingParams( + temperature=0.4, + top_p=0.9, + max_tokens=args.max_tokens, + repetition_penalty=1.05, + seed=SEED, + detokenize=True, + ) + sampling_params_list = [thinker_sampling_params] + + prompts = [query_result.inputs for _ in range(args.num_prompts)] + + if args.modalities is not None: + output_modalities = args.modalities.split(",") + for prompt in prompts: + prompt["modalities"] = output_modalities + + total_requests = len(prompts) + processed_count = 0 + print(f"Query type: {args.query_type}") + print(f"Number of prompts: {total_requests}") + + output_dir = args.output_dir + os.makedirs(output_dir, exist_ok=True) + + profiler_enabled = args.enable_profiler + if profiler_enabled: + omni.start_profile(stages=args.profiler_stages) + + for stage_outputs in omni.generate(prompts, sampling_params_list): + output = stage_outputs.request_output + if stage_outputs.final_output_type == "text": + request_id = output.request_id + text_output = output.outputs[0].text + lines = [] + lines.append("Prompt:\n") + lines.append(str(output.prompt) + "\n") + lines.append("Text Output:\n") + lines.append(str(text_output).strip() + "\n") + print(*lines, sep="") + + # Save to file + out_txt = os.path.join(output_dir, f"{request_id}.txt") + try: + with open(out_txt, "w", encoding="utf-8") as f: + f.writelines(lines) + print(f"Request ID: {request_id}, text saved to {out_txt}") + except Exception as e: + print(f"Failed to write output file {out_txt}: {e}") + + elif stage_outputs.final_output_type == "audio": + raise NotImplementedError("Add audio example after talker supported.") + + processed_count += 1 + if profiler_enabled and processed_count >= total_requests: + print(f"[Info] Processed {processed_count}/{total_requests}. Stopping profiler inside active loop...") + # Stop the profiler while workers are still alive + omni.stop_profile(stages=args.profiler_stages) + + print("[Info] Waiting 30s for workers to write trace files to disk...") + time.sleep(30) + print("[Info] Trace export wait time finished.") + + omni.close() + + +def parse_args(): + parser = FlexibleArgumentParser(description="Ming-flash-omni 2.0 offline inference example") + parser.add_argument( + "--query-type", + "-q", + type=str, + default="text", + choices=query_map.keys(), + help="Query type.", + ) + parser.add_argument( + "--stage-configs-path", + type=str, + default=None, + help="Path to a stage configs YAML file.", + ) + parser.add_argument( + "--log-stats", + action="store_true", + default=False, + help="Enable detailed statistics logging.", + ) + parser.add_argument("--init-timeout", type=int, default=2000, help="Timeout for initializing in seconds.") + parser.add_argument( + "--stage-init-timeout", + type=int, + default=2000, + help="Timeout for initializing a single stage in seconds.", + ) + parser.add_argument( + "--enable-profiler", + action="store_true", + default=False, + help="Enables profiling when set.", + ) + parser.add_argument( + "--profiler-stages", + type=int, + nargs="*", + default=[0], + help="List of stage IDs to profile. If not set, profiles all stages.", + ) + parser.add_argument( + "--image-path", + "-i", + type=str, + default=None, + help="Path to local image file. Uses default asset if not provided.", + ) + parser.add_argument( + "--audio-path", + "-a", + type=str, + default=None, + help="Path to local audio file. Uses default asset if not provided.", + ) + parser.add_argument( + "--video-path", + "-v", + type=str, + default=None, + help="Path to local video file. Uses default asset if not provided.", + ) + parser.add_argument( + "--num-frames", + type=int, + default=16, + help="Number of frames to extract from video.", + ) + parser.add_argument( + "--sampling-rate", + type=int, + default=16000, + help="Sampling rate for audio loading.", + ) + parser.add_argument( + "--max-tokens", + type=int, + default=16384, + help="Maximum tokens to generate.", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1, + help="Number of prompts to generate.", + ) + parser.add_argument( + "--modalities", + type=str, + default=None, + help="Output modalities (comma-separated).", + ) + parser.add_argument( + "--output-dir", + type=str, + default="output_ming", + help="Output directory for results.", + ) + + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + main(args) diff --git a/examples/online_serving/ming_flash_omni/README.md b/examples/online_serving/ming_flash_omni/README.md new file mode 100644 index 000000000..502232725 --- /dev/null +++ b/examples/online_serving/ming_flash_omni/README.md @@ -0,0 +1,204 @@ +# Ming-flash-omni 2.0 + +## Installation + +Please refer to [README.md](../../../README.md) + +## Run examples (Ming-flash-omni 2.0) + +### Launch the Server + +```bash +vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 +``` + +If you have custom stage configs file, launch the server with command below +```bash +vllm serve Jonathan1909/Ming-flash-omni-2.0 --omni --port 8091 --stage-configs-path /path/to/stage_configs_file +``` + +### Send Multi-modal Request + +#### Send request via python + +```bash +python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py --model Jonathan1909/Ming-flash-omni-2.0 --query-type use_mixed_modalities --port 8091 --host "localhost" --modalities text +``` + +The Python client supports the following command-line arguments: + +- `--query-type` (or `-q`): Query type. Options: `text`, `use_audio`, `use_image`, `use_video`, `use_mixed_modalities` +- `--video-path` (or `-v`): Path to local video file or URL. If not provided and query-type uses video, uses default video URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs. Example: `--video-path /path/to/video.mp4` or `--video-path https://example.com/video.mp4` +- `--image-path` (or `-i`): Path to local image file or URL. If not provided and query-type uses image, uses default image URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common image formats: JPEG, PNG, GIF, WebP. Example: `--image-path /path/to/image.jpg` or `--image-path https://example.com/image.png` +- `--audio-path` (or `-a`): Path to local audio file or URL. If not provided and query-type uses audio, uses default audio URL. Supports local file paths (automatically encoded to base64) or HTTP/HTTPS URLs and common audio formats: MP3, WAV, OGG, FLAC, M4A. Example: `--audio-path /path/to/audio.wav` or `--audio-path https://example.com/audio.mp3` +- `--prompt` (or `-p`): Custom text prompt/question. If not provided, uses default prompt for the selected query type. Example: `--prompt "What are the main activities shown in this video?"` +- `--modalities`: Output modalities. For now, only `text` is supported. Example: `--modalities text` + + +#### Send request via curl + +```bash +bash run_curl_multimodal_generation.sh text +bash run_curl_multimodal_generation.sh use_image +bash run_curl_multimodal_generation.sh use_audio +bash run_curl_multimodal_generation.sh use_video +bash run_curl_multimodal_generation.sh use_mixed_modalities +``` + +## Modality control + +Ming-flash-omni 2.0 currently supports text output only (thinker stage). + +| Modalities | Output | +|------------|--------| +| `["text"]` | Text only | +| Not specified | Text only (default) | + +### Using curl + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"} + ], + "modalities": ["text"] + }' +``` + +### Using OpenAI Python SDK + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") + +response = client.chat.completions.create( + model="Jonathan1909/Ming-flash-omni-2.0", + messages=[ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"}, + ], + modalities=["text"], +) +print(response.choices[0].message.content) +``` + +### Multi-modal input with OpenAI Python SDK + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") + +response = client.chat.completions.create( + model="Jonathan1909/Ming-flash-omni-2.0", + messages=[ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": "https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg"}}, + {"type": "text", "text": "Describe this image in detail."}, + ], + }, + ], + modalities=["text"], +) +print(response.choices[0].message.content) +``` + +## Streaming Output + +To enable streaming output: + +```bash +python examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py \ + --query-type use_image \ + --model Jonathan1909/Ming-flash-omni-2.0 \ + --modalities text \ + --stream +``` + +Or with the OpenAI Python SDK: + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") + +response = client.chat.completions.create( + model="Jonathan1909/Ming-flash-omni-2.0", + messages=[ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"}, + ], + modalities=["text"], + stream=True, +) +for chunk in response: + for choice in chunk.choices: + if hasattr(choice, "delta") and choice.delta.content: + print(choice.delta.content, end="", flush=True) +print() +``` + +Or using curl: + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking off"}]}, + {"role": "user", "content": "请详细介绍鹦鹉的生活习性。"} + ], + "modalities": ["text"], + "stream": true, + }' +``` + + +## Reasoning (Thinking Mode) + +To enable reasoning/thinking mode, change `detailed thinking off` to `detailed thinking on` in the system prompt: + +### Using curl + +```bash +curl http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d '{ + "model": "Jonathan1909/Ming-flash-omni-2.0", + "messages": [ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking on"}]}, + {"role": "user", "content": [ + {"type": "image_url", "image_url": {"url": "https://example.com/math_problem.png"}}, + {"type": "text", "text": "Solve this math problem step by step."} + ]} + ], + "modalities": ["text"] + }' +``` + +### Using OpenAI Python SDK + +```python +from openai import OpenAI + +client = OpenAI(base_url="http://localhost:8091/v1", api_key="EMPTY") + +response = client.chat.completions.create( + model="Jonathan1909/Ming-flash-omni-2.0", + messages=[ + {"role": "system", "content": [{"type": "text", "text": "你是一个友好的AI助手。\n\ndetailed thinking on"}]}, + {"role": "user", "content": "If a train travels 120 km in 2 hours, what is its average speed?"}, + ], + modalities=["text"], +) +print(response.choices[0].message.content) +``` diff --git a/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh b/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh new file mode 100755 index 000000000..768a424e4 --- /dev/null +++ b/examples/online_serving/ming_flash_omni/run_curl_multimodal_generation.sh @@ -0,0 +1,145 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Server port +PORT="${PORT:-8091}" +# Default query type +QUERY_TYPE="${1:-text}" + +# Validate query type +if [[ ! "$QUERY_TYPE" =~ ^(text|use_audio|use_image|use_video|use_mixed_modalities)$ ]]; then + echo "Error: Invalid query type '$QUERY_TYPE'" + echo "Usage: $0 [text|use_audio|use_image|use_video|use_mixed_modalities]" + echo " text: Text-only query" + echo " use_audio: Audio + Text query" + echo " use_image: Image + Text query" + echo " use_video: Video + Text query" + echo " use_mixed_modalities: Audio + Image + Video + Text query" + exit 1 +fi + +thinker_sampling_params='{ + "temperature": 0.4, + "top_p": 0.9, + "top_k": -1, + "max_tokens": 16384, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.05 +}' +# Above is optional, it has a default setting in stage_configs of the corresponding model. + +# Define URLs for assets +MARY_HAD_LAMB_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/mary_had_lamb.ogg" +CHERRY_BLOSSOM_IMAGE_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/vision_model_images/cherry_blossom.jpg" +SAMPLE_VIDEO_URL="https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4" + +# Build user content based on query type +case "$QUERY_TYPE" in + text) + user_content='[ + { + "type": "text", + "text": "请详细介绍鹦鹉的生活习性。" + } + ]' + ;; + use_image) + user_content='[ + { + "type": "image_url", + "image_url": { + "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'" + } + }, + { + "type": "text", + "text": "Describe this image in detail." + } + ]' + ;; + use_audio) + user_content='[ + { + "type": "audio_url", + "audio_url": { + "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'" + } + }, + { + "type": "text", + "text": "Please recognize the language of this speech and transcribe it. Format: oral." + } + ]' + ;; + use_video) + user_content='[ + { + "type": "video_url", + "video_url": { + "url": "'"$SAMPLE_VIDEO_URL"'" + } + }, + { + "type": "text", + "text": "Describe what is happening in this video." + } + ]' + ;; + use_mixed_modalities) + user_content='[ + { + "type": "image_url", + "image_url": { + "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'" + } + }, + { + "type": "audio_url", + "audio_url": { + "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'" + } + }, + { + "type": "text", + "text": "Describe the image, and recognize the language of this speech and transcribe it. Format: oral" + } + ]' + ;; +esac + +echo "Running query type: $QUERY_TYPE" +echo "" + +request_body=$(cat < str: + """Build a Ming chat prompt.""" + return ( + f"SYSTEM{SYSTEM_PROMPT}{EOS_TOKEN}HUMAN{user_text}{EOS_TOKEN}ASSISTANT" + ) + + +def get_eager_config(): + path = modify_stage_config( + str(Path(__file__).parent.parent / "stage_configs" / "bailingmm_moe_v2_lite_ci.yaml"), + updates={ + "stage_args": { + 0: { + "engine_args.enforce_eager": "true", + }, + }, + }, + ) + return path + + +stage_configs = [get_eager_config()] +test_params = [(model, stage_config) for model in models for stage_config in stage_configs] + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", test_params, indirect=True) +def test_text_to_text(omni_runner, omni_runner_handler) -> None: + """ + Test text-only input processing and text output generation. + Input Modal: text + Output Modal: text + """ + prompt = build_prompt("请详细介绍鹦鹉的生活习性。") + request_config = {"prompts": prompt, "modalities": ["text"]} + + omni_runner_handler.send_request(request_config) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", test_params, indirect=True) +def test_image_to_text(omni_runner, omni_runner_handler) -> None: + """ + Test image understanding with text output. + Input Modal: image + text + Output Modal: text + """ + image = generate_synthetic_image(224, 224)["np_array"] + prompt = build_prompt(f"{IMAGE_TOKEN}Describe this image briefly.") + request_config = {"prompts": prompt, "images": image, "modalities": ["text"]} + + omni_runner_handler.send_request(request_config) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", test_params, indirect=True) +def test_audio_to_text(omni_runner, omni_runner_handler) -> None: + """ + Test audio understanding with text output. + Input Modal: audio + text + Output Modal: text + """ + audio = generate_synthetic_audio(2, 1, 16000)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + prompt = build_prompt(f"{AUDIO_TOKEN}Please recognize the language of this speech and transcribe it. Format: oral.") + request_config = {"prompts": prompt, "audios": audio, "modalities": ["text"]} + + omni_runner_handler.send_request(request_config) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", test_params, indirect=True) +def test_video_to_text(omni_runner, omni_runner_handler) -> None: + """ + Test video understanding with text output. + Input Modal: video + text + Output Modal: text + """ + video = generate_synthetic_video(224, 224, 30)["np_array"] + prompt = build_prompt(f"{VIDEO_TOKEN}Describe what is happening in this video.") + request_config = {"prompts": prompt, "videos": video, "modalities": ["text"]} + + omni_runner_handler.send_request(request_config) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_runner", test_params, indirect=True) +def test_mixed_to_text(omni_runner, omni_runner_handler) -> None: + """ + Test mixed modality input (image + audio) with text output. + Input Modal: image + audio + text + Output Modal: text + """ + image = generate_synthetic_image(224, 224)["np_array"] + audio = generate_synthetic_audio(2, 1, 16000)["np_array"] + if len(audio.shape) == 2: + audio = audio.squeeze() + prompt = build_prompt(f"{IMAGE_TOKEN}{AUDIO_TOKEN}Describe the image and transcribe the audio.") + request_config = {"prompts": prompt, "images": image, "audios": audio, "modalities": ["text"]} + + omni_runner_handler.send_request(request_config) diff --git a/tests/e2e/online_serving/test_ming_flash_omni.py b/tests/e2e/online_serving/test_ming_flash_omni.py new file mode 100644 index 000000000..35b7b64c0 --- /dev/null +++ b/tests/e2e/online_serving/test_ming_flash_omni.py @@ -0,0 +1,247 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +E2E online serving tests for Ming-flash-omni-2.0 model (Thinker stage). +Tests multimodal understanding via OpenAI-compatible API. +""" + +import os +from pathlib import Path + +import pytest + +from tests.conftest import ( + OmniServerParams, + dummy_messages_from_mix_data, + generate_synthetic_audio, + generate_synthetic_image, + generate_synthetic_video, + modify_stage_config, +) +from tests.utils import hardware_test + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "0" + +models = ["Jonathan1909/Ming-flash-omni-2.0"] + + +def get_eager_config(): + path = modify_stage_config( + str(Path(__file__).parent.parent / "stage_configs" / "bailingmm_moe_v2_lite_ci.yaml"), + updates={ + "stage_args": { + 0: { + "engine_args.enforce_eager": "true", + }, + }, + }, + ) + return path + + +stage_configs = [get_eager_config()] + +# Create parameter combinations for model and stage config +test_params = [ + OmniServerParams(model=model, stage_config_path=stage_config) for model in models for stage_config in stage_configs +] + + +def get_system_prompt(): + return { + "role": "system", + "content": [ + { + "type": "text", + "text": "你是一个友好的AI助手。\n\ndetailed thinking off", + } + ], + } + + +def get_prompt(prompt_type="text_only"): + prompts = { + "text_only": "What is the capital of China? Answer in 20 words.", + "text_image": "What is in this image?", + "text_audio": "What is in this audio?", + "text_video": "What is in this video?", + "mix": "What is recited in the audio? What is in this image? What is in this video?", + } + return prompts.get(prompt_type, prompts["text_only"]) + + +def get_max_batch_size(size_type="few"): + batch_sizes = {"few": 5, "medium": 100, "large": 256} + return batch_sizes.get(size_type, 5) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_to_text_001(omni_server, openai_client) -> None: + """ + Input Modal: text + Output Modal: text + Input Setting: stream=False + Datasets: single request + """ + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + content_text=get_prompt("text_only"), + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": False, + "modalities": ["text"], + "key_words": {"text": ["beijing"]}, + } + + openai_client.send_omni_request(request_config) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_text_to_text_stream_001(omni_server, openai_client) -> None: + """ + Input Modal: text + Output Modal: text + Input Setting: stream=True + Datasets: few requests + """ + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + content_text=get_prompt("text_only"), + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": True, + "modalities": ["text"], + "key_words": {"text": ["beijing"]}, + } + + openai_client.send_omni_request(request_config, request_num=get_max_batch_size()) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_image_to_text_001(omni_server, openai_client) -> None: + """ + Input Modal: image + text + Output Modal: text + Input Setting: stream=True + Datasets: single request + """ + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + image_data_url=image_data_url, + content_text=get_prompt("text_image"), + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": True, + "modalities": ["text"], + } + + openai_client.send_omni_request(request_config) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_audio_to_text_001(omni_server, openai_client) -> None: + """ + Input Modal: audio + text + Output Modal: text + Input Setting: stream=True + Datasets: single request + """ + audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(2, 1)['base64']}" + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + audio_data_url=audio_data_url, + content_text=get_prompt("text_audio"), + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": True, + "modalities": ["text"], + } + + openai_client.send_omni_request(request_config) + + +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_video_to_text_001(omni_server, openai_client) -> None: + """ + Input Modal: video + text + Output Modal: text + Input Setting: stream=False + Datasets: single request + """ + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}" + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + video_data_url=video_data_url, + content_text=get_prompt("text_video"), + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": False, + "modalities": ["text"], + } + + openai_client.send_omni_request(request_config) + + +@pytest.mark.advanced_model +@pytest.mark.core_model +@pytest.mark.omni +@hardware_test(res={"cuda": "H100"}, num_cards=4) +@pytest.mark.parametrize("omni_server", test_params, indirect=True) +def test_mix_to_text_001(omni_server, openai_client) -> None: + """ + Input Modal: text + audio + image + video + Output Modal: text + Input Setting: stream=True + Datasets: single request + """ + video_data_url = f"data:video/mp4;base64,{generate_synthetic_video(224, 224, 300)['base64']}" + image_data_url = f"data:image/jpeg;base64,{generate_synthetic_image(224, 224)['base64']}" + audio_data_url = f"data:audio/wav;base64,{generate_synthetic_audio(2, 1)['base64']}" + messages = dummy_messages_from_mix_data( + system_prompt=get_system_prompt(), + video_data_url=video_data_url, + image_data_url=image_data_url, + audio_data_url=audio_data_url, + content_text=get_prompt("mix"), + ) + + request_config = { + "model": omni_server.model, + "messages": messages, + "stream": True, + "modalities": ["text"], + } + + openai_client.send_omni_request(request_config) diff --git a/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml b/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml new file mode 100644 index 000000000..fb0c72cc5 --- /dev/null +++ b/tests/e2e/stage_configs/bailingmm_moe_v2_lite_ci.yaml @@ -0,0 +1,35 @@ +# Thinker stage only +stage_args: + - stage_id: 0 + stage_type: llm + runtime: + devices: "0,1,2,3" + max_batch_size: 1 + engine_args: + model_stage: thinker + model_arch: MingFlashOmniForConditionalGeneration + worker_type: ar + scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler + gpu_memory_utilization: 0.9 + enforce_eager: false + trust_remote_code: true + engine_output_type: latent + distributed_executor_backend: "mp" + enable_prefix_caching: false + max_num_batched_tokens: 32768 + max_model_len: 32768 + tensor_parallel_size: 4 + hf_config_name: llm_config + load_format: dummy + mm_processor_cache_gb: 0 + final_output: true + final_output_type: text + is_comprehension: true + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + max_tokens: 100 + repetition_penalty: 1.05 + seed: 42 + detokenize: true + ignore_eos: false diff --git a/vllm_omni/model_executor/models/ming_flash_omni/__init__.py b/vllm_omni/model_executor/models/ming_flash_omni/__init__.py new file mode 100644 index 000000000..d7fa44fd7 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/__init__.py @@ -0,0 +1,18 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. + +from .ming_flash_omni import MingFlashOmniForConditionalGeneration +from .ming_flash_omni_thinker import ( + MingFlashOmniThinkerDummyInputsBuilder, + MingFlashOmniThinkerForConditionalGeneration, + MingFlashOmniThinkerMultiModalProcessor, + MingFlashOmniThinkerProcessingInfo, +) + +__all__ = [ + "MingFlashOmniForConditionalGeneration", + "MingFlashOmniThinkerForConditionalGeneration", + "MingFlashOmniThinkerProcessingInfo", + "MingFlashOmniThinkerMultiModalProcessor", + "MingFlashOmniThinkerDummyInputsBuilder", +] diff --git a/vllm_omni/model_executor/models/ming_flash_omni/audio_encoder.py b/vllm_omni/model_executor/models/ming_flash_omni/audio_encoder.py new file mode 100644 index 000000000..6ca199011 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/audio_encoder.py @@ -0,0 +1,246 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright 2024 ANT Group and the HuggingFace Inc. team. +# Copyright (c) 2022 OpenAI +# Adapted from Ming repository modeling_whisper_encoder.py +# https://github.com/inclusionAI/Ming + +import operator +from collections.abc import Iterable +from itertools import accumulate + +import torch +import torch.nn as nn +import torch.nn.functional as F +from vllm.logger import init_logger +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.backends.utils.fa import HAS_FLASH_ATTN, flash_attn_varlen_func +from vllm_omni.model_executor.models.whisper_utils import Conv1d, Linear, sinusoids + +logger = init_logger(__name__) + + +class MultiHeadAttention(nn.Module): + """Multi-head attention with packed sequence support. + Adapted from Qwen3-TTS WhisperEncoder. + """ + + def __init__(self, n_state: int, n_head: int, use_flash_attn: bool = True): + super().__init__() + self.n_head = n_head + self.query = Linear(n_state, n_state) + self.key = Linear(n_state, n_state, bias=False) + self.value = Linear(n_state, n_state) + self.out = Linear(n_state, n_state) + + if use_flash_attn and not HAS_FLASH_ATTN: + logger.warning("flash-attn is not available. Fallback to manual PyTorch version") + self.use_flash_attn = use_flash_attn and HAS_FLASH_ATTN + + def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: + """Forward pass with packed sequence support. + + Args: + x: [total_tokens, n_state] packed sequence + cu_seqlens: [num_seqs + 1] cumulative sequence lengths, e.g. [0, len1, len1+len2, ...] + + Returns: + [total_tokens, n_state] attention output + """ + q = self.query(x) + k = self.key(x) + v = self.value(x) + + n_ctx, n_state = q.shape + head_dim = n_state // self.n_head + + q = q.view(n_ctx, self.n_head, head_dim) + k = k.view(n_ctx, self.n_head, head_dim) + v = v.view(n_ctx, self.n_head, head_dim) + + # Try flash attention varlen + if self.use_flash_attn and cu_seqlens is not None and q.dtype in [torch.float16, torch.bfloat16]: + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + attn_output = flash_attn_varlen_func(q, k, v, cu_seqlens, cu_seqlens, max_seqlen, max_seqlen) + else: + attn_output = self._manual_attention(q, k, v, cu_seqlens) + + # Reshape back: [T, H, D] -> [T, H*D] + attn_output = attn_output.contiguous().view(n_ctx, n_state) + return self.out(attn_output) + + def _manual_attention( + self, q: torch.Tensor, k: torch.Tensor, v: torch.Tensor, cu_seqlens: torch.Tensor + ) -> torch.Tensor: + """Manual attention for variable-length sequences (fallback).""" + _, n_head, head_dim = q.shape + scale = head_dim**-0.5 + + # Unpack sequences and pad to max length + seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() + batch_size = len(seqlens) + max_seqlen = max(seqlens) + + # Create padded tensors + q_padded = torch.zeros(batch_size, max_seqlen, n_head, head_dim, dtype=q.dtype, device=q.device) + k_padded = torch.zeros_like(q_padded) + v_padded = torch.zeros_like(q_padded) + + # Fill with actual sequences + for i in range(batch_size): + start_idx = cu_seqlens[i] + end_idx = cu_seqlens[i + 1] + seq_len = seqlens[i] + q_padded[i, :seq_len] = q[start_idx:end_idx] + k_padded[i, :seq_len] = k[start_idx:end_idx] + v_padded[i, :seq_len] = v[start_idx:end_idx] + + # Transpose for attention: [B, H, T, D] + q_padded = q_padded.transpose(1, 2) + k_padded = k_padded.transpose(1, 2) + v_padded = v_padded.transpose(1, 2) + + # Create attention mask for variable lengths: 0 for valid positions, -inf for padding + padding_mask = ( + torch.arange(max_seqlen, device=q.device)[None, :] >= torch.tensor(seqlens, device=q.device)[:, None] + ) + attn_mask = torch.zeros(batch_size, 1, 1, max_seqlen, dtype=q.dtype, device=q.device) + attn_mask = attn_mask.masked_fill(padding_mask.unsqueeze(1).unsqueeze(2), -torch.finfo(q.dtype).max) + + # Compute attention + attn_scores = torch.matmul(q_padded, k_padded.transpose(-2, -1)) * scale + attn_scores = attn_scores + attn_mask + attn_weights = F.softmax(attn_scores, dim=-1) + context = torch.matmul(attn_weights, v_padded) + + # Transpose back: [B, H, T, D] -> [B, T, H, D] + context = context.transpose(1, 2).contiguous() + output_packed = torch.cat([context[i, : seqlens[i]] for i in range(batch_size)], dim=0) + + return output_packed + + +class ResidualAttentionBlock(nn.Module): + """Whisper-style residual attention block with packed sequence support. + + Adapted from + https://github.com/openai/whisper/blob/v20250625/whisper/model.py + vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py + """ + + def __init__(self, n_state: int, n_head: int, use_flash_attn: bool = True): + super().__init__() + self.attn = MultiHeadAttention(n_state, n_head, use_flash_attn=use_flash_attn) + self.attn_ln = nn.LayerNorm(n_state) + + n_mlp = n_state * 4 + self.mlp = nn.Sequential( + Linear(n_state, n_mlp), + nn.GELU(), + Linear(n_mlp, n_state), + ) + self.mlp_ln = nn.LayerNorm(n_state) + + def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor) -> torch.Tensor: + x = x + self.attn(self.attn_ln(x), cu_seqlens=cu_seqlens) + x = x + self.mlp(self.mlp_ln(x)) + return x + + +class WhisperAudioEncoder(nn.Module): + """Whisper audio encoder for Ming with packed sequence support. + + Adapted from + https://github.com/openai/whisper/blob/v20250625/whisper/model.py + vllm_omni/model_executor/models/qwen3_tts/tokenizer_25hz/vq/whisper_encoder.py + """ + + def __init__( + self, + n_mels: int = 128, + n_ctx: int = 15000, + n_state: int = 1280, + n_head: int = 20, + n_layer: int = 32, + use_flash_attn: bool = True, + ): + super().__init__() + self.conv1 = Conv1d(n_mels, n_state, kernel_size=3, padding=1) + self.conv2 = Conv1d(n_state, n_state, kernel_size=3, stride=2, padding=1) + # self.positional_embedding = nn.Parameter(torch.empty(n_ctx, n_state)) + self.register_buffer("positional_embedding", sinusoids(n_ctx, n_state)) + self.blocks = nn.ModuleList( + [ResidualAttentionBlock(n_state, n_head, use_flash_attn=use_flash_attn) for _ in range(n_layer)] + ) + self.ln_post = nn.LayerNorm(n_state) + self.audio_emb_dim = n_state + + self.n_layer = n_layer + self.n_mels = n_mels + self.use_flash_attn = use_flash_attn + + def forward( + self, + x_list: list[torch.Tensor], + audio_lens: list[int], + ) -> torch.Tensor: + """Forward pass with packed sequence format for variable-length inputs. + + Args: + x_list: List of [n_mels, T_i] mel spectrogram features for each audio + audio_lens: List of original audio lengths in frames + + Returns: + [total_T', n_state] packed encoded audio features, where + total_T' is the sum of all encoded sequence lengths + """ + # Cast inputs to model dtype + target_dtype = self.conv1.weight.dtype + x_list = [x.to(target_dtype) for x in x_list] + + encoded_list = [] + encoded_lens = [] + for mel_spec in x_list: + # mel_spec: [n_mels, T] - process through conv layers + x = mel_spec.unsqueeze(0) # [1, n_mels, T] + x = F.gelu(self.conv1(x)) + x = F.gelu(self.conv2(x)) + x = x.squeeze(0).transpose(0, 1) # [T', n_state] + + # Add positional embedding + seq_len = x.shape[0] + positional_embedding = self.positional_embedding[:seq_len, :] + x = (x + positional_embedding).to(x.dtype) + + encoded_list.append(x) + encoded_lens.append(seq_len) + + x_packed = torch.cat(encoded_list, dim=0) # [total_T', n_state] + + cu_seqlens = list(accumulate(encoded_lens, func=operator.add, initial=0)) + cu_seqlens = torch.tensor(cu_seqlens, device=x_packed.device, dtype=torch.int32) + + for block in self.blocks: + x_packed = block(x_packed, cu_seqlens=cu_seqlens) + + x_packed = self.ln_post(x_packed) + return x_packed + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + params_dict: dict[str, torch.Tensor] = { + **dict(self.named_parameters(remove_duplicate=False)), + **dict(self.named_buffers()), + } + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + if name not in params_dict: + logger.warning("Skipping unknown audio encoder weight: %s", name) + continue + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + loaded_params.add(name) + + return loaded_params diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py new file mode 100644 index 000000000..87728890b --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright 2024 ANT Group and the HuggingFace Inc. team. All rights reserved. +# Adapted from Ming repository modeling_bailingmm2.py +# https://github.com/inclusionAI/Ming +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. + +"""Ming-flash-omni-2.0 unified model (thinker + imagegen + talker).""" + +from collections.abc import Iterable + +import torch +import torch.nn as nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import ( + SupportsMRoPE, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.module_mapping import MultiModelKeys +from vllm.model_executor.models.utils import ( + init_vllm_registered_model, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors + +from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights +from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config, MingFlashOmniConfig + +from .ming_flash_omni_thinker import ( + MingFlashOmniThinkerDummyInputsBuilder, + MingFlashOmniThinkerMultiModalProcessor, + MingFlashOmniThinkerProcessingInfo, +) + +logger = init_logger(__name__) + + +@MULTIMODAL_REGISTRY.register_processor( + MingFlashOmniThinkerMultiModalProcessor, + info=MingFlashOmniThinkerProcessingInfo, + dummy_inputs=MingFlashOmniThinkerDummyInputsBuilder, +) +class MingFlashOmniForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + SupportsMRoPE, + CustomProcessMixin, +): + """Unified Ming-flash-omni-2.0 model combining thinker, imagegen, and talker.""" + + supports_multimodal = True + requires_raw_input_tokens: bool = True + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.have_multimodal_outputs = True + self.has_preprocess = False + self.has_postprocess = False + + config = vllm_config.model_config.hf_config + + self.vllm_config = vllm_config + self.config = config + + if isinstance(config, MingFlashOmniConfig): + thinker_config = config.thinker_config + else: + thinker_config = config + + self.thinker_config: BailingMM2Config = thinker_config + self.model_stage = vllm_config.model_config.model_stage + + if self.model_stage == "thinker": + thinker_vllm_config = vllm_config.with_hf_config( + thinker_config, architectures=["MingFlashOmniThinkerForConditionalGeneration"] + ) + self.thinker = init_vllm_registered_model( + vllm_config=thinker_vllm_config, + prefix=maybe_prefix(prefix, "thinker"), + architectures=["MingFlashOmniThinkerForConditionalGeneration"], + ) + self.model = self.thinker + self.imagegen = None + self.talker = None + + elif self.model_stage == "imagegen": + # TODO: Implement image generator stage + raise NotImplementedError( + "Image generation stage is not yet implemented. Please use model_stage='thinker' for now." + ) + + elif self.model_stage == "talker": + # TODO: Implement talker (TTS) stage + raise NotImplementedError( + "Talker (TTS) stage is not yet implemented. Please use model_stage='thinker' for now." + ) + + else: + raise ValueError( + f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'imagegen', 'talker'" + ) + + # Set up intermediate tensors + self.make_empty_intermediate_tensors = ( + self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None + ) + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **kwargs, + ) -> OmniOutput: + return self.model.forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **kwargs, + ) + + def compute_logits( + self, + hidden_states: torch.Tensor, + sampling_metadata=None, + ) -> torch.Tensor | None: + if hasattr(self.model, "compute_logits"): + return self.model.compute_logits(hidden_states, sampling_metadata) + return None + + def sample( + self, + logits: torch.Tensor, + sampling_metadata, + ): + if hasattr(self.model, "sample"): + return self.model.sample(logits, sampling_metadata) + raise NotImplementedError("sample method not available on current stage") + + def get_mrope_input_positions(self, *args, **kwargs): + if hasattr(self.model, "get_mrope_input_positions"): + return self.model.get_mrope_input_positions(*args, **kwargs) + raise NotImplementedError("get_mrope_input_positions not available on current stage") + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loaded_weights = set() + thinker_weights = [] + imagegen_weights = [] + talker_weights = [] + + for name, value in weights: + if name.startswith("thinker."): + thinker_weights.append((name, value)) + elif name.startswith("imagegen."): + imagegen_weights.append((name, value)) + elif name.startswith("talker."): + talker_weights.append((name, value)) + else: + # Weights without prefix go to thinker by default + thinker_weights.append((name, value)) + + if self.model_stage == "thinker" and thinker_weights: + # Remove "thinker." prefix before loading + thinker_weights_stripped = [ + (name.replace("thinker.", "", 1) if name.startswith("thinker.") else name, value) + for name, value in thinker_weights + ] + thinker_loaded = self.thinker.load_weights(thinker_weights_stripped) + thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker") + loaded_weights.update(thinker_loaded) + + # TODO: Load imagegen weights when implemented + # TODO: Load talker weights when implemented + + return loaded_weights + + def get_mm_mapping(self) -> MultiModelKeys: + return MultiModelKeys.from_string_field( + language_model="thinker.language_model", + connector=["thinker.linear_proj.", "thinker.linear_proj_audio."], + tower_model=["thinker.vision.", "thinker.audio."], + ) + + @property + def sampler(self): + if hasattr(self.model, "sampler"): + return self.model.sampler + return None + + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + *, + is_multimodal=None, + ) -> torch.Tensor: + return self.model.embed_input_ids( + input_ids, + multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + def embed_multimodal(self, **kwargs): + return self.model.embed_multimodal(**kwargs) diff --git a/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py new file mode 100644 index 000000000..bde7477b9 --- /dev/null +++ b/vllm_omni/model_executor/models/ming_flash_omni/ming_flash_omni_thinker.py @@ -0,0 +1,893 @@ +# SPDX-License-Identifier: Apache-2.0 +# Copyright 2025 The vLLM-Omni team. +# Copyright 2024 ANT Group and the HuggingFace Inc. team. +# Adapted from Ming repository modeling_bailingmm2.py and processing_bailingmm2.py +# https://github.com/inclusionAI/Ming + +"""Ming-flash-omni-2.0 Thinker stage implementation (multimodal understanding).""" + +from collections.abc import Iterable, Iterator, Mapping, Sequence +from typing import Annotated, Any + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.feature_extraction_utils import BatchFeature +from vllm.config import VllmConfig +from vllm.config.multimodal import BaseDummyOptions +from vllm.inputs import MultiModalDataDict +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsMRoPE, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.qwen2_5_vl import ( + Qwen2_5_VLImageInputs, + Qwen2_5_VLImagePixelInputs, + Qwen2_5_VLVideoInputs, + Qwen2_5_VLVideoPixelInputs, +) +from vllm.model_executor.models.qwen2_vl import ( + Qwen2VLProcessingInfo, +) +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + _merge_multimodal_embeddings, + maybe_prefix, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalFeatureSpec, + MultiModalFieldConfig, + MultiModalKwargsItems, +) +from vllm.multimodal.parse import ( + AudioProcessorItems, + ImageProcessorItems, + MultiModalDataItems, + MultiModalDataParser, + VideoProcessorItems, +) +from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, + BaseMultiModalProcessor, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.sequence import IntermediateTensors +from vllm.utils.tensor_schema import TensorSchema, TensorShape + +from vllm_omni.model_executor.custom_process_mixin import CustomProcessMixin +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.transformers_utils.configs.ming_flash_omni import BailingMM2Config +from vllm_omni.transformers_utils.processors.ming import ( + PLACEHOLDER_AUDIO_TOKEN_IN_TEXT, + PLACEHOLDER_IMAGE_TOKEN_IN_TEXT, + PLACEHOLDER_VIDEO_TOKEN_IN_TEXT, + MingFlashOmniProcessor, + MingWhisperFeatureExtractor, +) + +from .audio_encoder import WhisperAudioEncoder +from .modeling_bailing_moe_v2 import BailingMoeV2ForCausalLM +from .projectors import AudioProjector, VisionProjector +from .vision_encoder import MingVisionEncoder + +logger = init_logger(__name__) + + +class MingAudioInput(TensorSchema): + """ + Dimensions: + - b: Batch size + - l: Total audio frames (clips concatenated along the time axis) + - nm: Number of mel bins + - N: Max number of audio clips per batch item + """ + + audio_feats: Annotated[ + torch.Tensor, + TensorShape("b", "l", "nm"), + ] + + audio_feats_lengths: Annotated[ + torch.Tensor, + TensorShape("b", "N"), + ] + + +class MingFlashOmniThinkerProcessingInfo(Qwen2VLProcessingInfo): + def get_hf_config(self) -> BailingMM2Config: + return self.ctx.get_hf_config(BailingMM2Config) + + def get_hf_processor(self, **kwargs: object): + return self.ctx.get_hf_processor(MingFlashOmniProcessor, **kwargs) + + def get_target_channels(self) -> int: + # See `_normalize_audio_tensor` in vllm_omni/transformers_utils/processors/ming.py + return 1 + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"image": None, "video": None, "audio": None} + + def get_mm_max_tokens_per_item( + self, + seq_len: int, + mm_counts: Mapping[str, int], + ) -> Mapping[str, int]: + mm_counts = mm_counts or {} + requested_modalities = {m for m, c in mm_counts.items() if c > 0} + mm_max_tokens: dict[str, int] = {} + + if requested_modalities & {"image", "video"}: + vl_tokens = super().get_mm_max_tokens_per_item( + seq_len=seq_len, + mm_counts=mm_counts, + ) + mm_max_tokens.update({m: vl_tokens[m] for m in ["image", "video"] if m in requested_modalities}) + + if "audio" in requested_modalities: + # TODO: consider computing from audio config + mm_max_tokens["audio"] = 3000 + + return mm_max_tokens + + def get_feature_extractor(self, **kwargs: object) -> MingWhisperFeatureExtractor: + hf_processor = self.get_hf_processor(**kwargs) + feature_extractor = hf_processor.audio_processor + assert isinstance(feature_extractor, MingWhisperFeatureExtractor) + return feature_extractor + + def get_data_parser(self): + feature_extractor = self.get_feature_extractor() + return MultiModalDataParser( + target_sr=feature_extractor.sampling_rate, + target_channels=self.get_target_channels(), + expected_hidden_size=self._get_expected_hidden_size(), + ) + + +class MingFlashOmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[MingFlashOmniThinkerProcessingInfo]): + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: + num_images = mm_counts.get("image", 0) + num_videos = mm_counts.get("video", 0) + num_audios = mm_counts.get("audio", 0) + + hf_processor = self.info.get_hf_processor() + + audio_token: str = hf_processor.audio_token + image_token: str = hf_processor.image_token + video_token: str = hf_processor.video_token + + return image_token * num_images + video_token * num_videos + audio_token * num_audios + + def get_dummy_mm_data( + self, + seq_len: int, + mm_counts: Mapping[str, int], + mm_options: Mapping[str, BaseDummyOptions] | None = None, + ) -> MultiModalDataDict: + num_images = mm_counts.get("image", 0) + num_videos = mm_counts.get("video", 0) + num_audios = mm_counts.get("audio", 0) + + # Default dimensions for dummy data + image_width, image_height = 448, 448 + video_width, video_height = 448, 448 + num_frames = 8 + audio_duration = 3.0 # seconds + sample_rate = 16000 + + audio_length = int(audio_duration * sample_rate) + + mm_data: MultiModalDataDict = { + "image": self._get_dummy_images( + width=image_width, + height=image_height, + num_images=num_images, + ), + "video": self._get_dummy_videos( + width=video_width, + height=video_height, + num_frames=num_frames, + num_videos=num_videos, + ), + "audio": [(np.random.randn(audio_length).astype(np.float32), sample_rate) for _ in range(num_audios)], + } + + return mm_data + + +class MingFlashOmniThinkerMultiModalProcessor(BaseMultiModalProcessor[MingFlashOmniThinkerProcessingInfo]): + """Multimodal processor for Ming-flash-omni Thinker stage. + + Handles preprocessing of 1) image, 2) video, and 3) audio inputs, + and expands placeholder tokens to the correct number of patch tokens. + """ + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + tokenizer = self.info.get_tokenizer() + # might want to add a fallback to resolve token ids + # vocab = tokenizer.get_vocab() + thinker_config = self.info.get_hf_config() + + # patch/delimiter token IDs (used in replacement sequences) + image_start_token_id = thinker_config.llm_config.image_start_token + image_patch_token_id = thinker_config.llm_config.image_patch_token + image_end_token_id = thinker_config.llm_config.image_end_token + + video_start_token_id = thinker_config.llm_config.video_start_token + frame_patch_token_id = thinker_config.llm_config.video_patch_token + video_end_token_id = thinker_config.llm_config.video_end_token + + audio_start_token_id = thinker_config.llm_config.audio_start_token + audio_patch_token_id = thinker_config.llm_config.audio_patch_token + audio_end_token_id = thinker_config.llm_config.audio_end_token + + vision_config = thinker_config.vision_config + spatial_merge_size = vision_config.spatial_merge_size if vision_config else 2 + + newline_token_ids: list[int] = tokenizer.encode("\n", add_special_tokens=False) + + out_mm_data = out_mm_kwargs.get_data() + + def get_replacement_image(item_idx: int) -> PromptUpdateDetails: + """Generate token sequence for an image.""" + grid_thw = out_mm_data.get("image_grid_thw") + if grid_thw is None: + raise ValueError( + "image_grid_thw missing from processor output; " + "cannot determine image patch count for prompt replacement." + ) + if isinstance(grid_thw, torch.Tensor): + thw = grid_thw[item_idx] + num_patches = int(thw.prod().item()) // (spatial_merge_size**2) + else: + thw = grid_thw[item_idx] + num_patches = (thw[0] * thw[1] * thw[2]) // (spatial_merge_size**2) + + # Build token sequence: *N \n + # the newline token is added in purpose from original model processing + tokens: list[int] = [] + tokens.append(image_start_token_id) + tokens.extend([image_patch_token_id] * num_patches) + tokens.append(image_end_token_id) + # Refer to Ming's BailingMM2Processor._expand_image_tokens + # https://github.com/inclusionAI/Ming/blob/3954fcb880ff5e61ff128bcf7f1ec344d46a6fe3/processing_bailingmm2.py + tokens.extend(newline_token_ids) + + # Only tokens receive multimodal embeddings + return PromptUpdateDetails.select_token_id(tokens, image_patch_token_id) + + def get_replacement_video(item_idx: int) -> PromptUpdateDetails: + """Generate token sequence for a video.""" + grid_thw = out_mm_data.get("video_grid_thw", None) + if grid_thw is None: + raise ValueError( + "video_grid_thw missing from processor output; " + "cannot determine video patch count for prompt replacement." + ) + if isinstance(grid_thw, torch.Tensor): + thw = grid_thw[item_idx] + num_patches = int(thw.prod().item()) // (spatial_merge_size**2) + else: + thw = grid_thw[item_idx] + num_patches = (thw[0] * thw[1] * thw[2]) // (spatial_merge_size**2) + + # Build token sequence: \n + # the newline token is added in purpose from original model processing + tokens: list[int] = [] + tokens.append(video_start_token_id) + tokens.extend([frame_patch_token_id] * num_patches) + tokens.append(video_end_token_id) + tokens.extend(newline_token_ids) + + # Only tokens receive multimodal embeddings + return PromptUpdateDetails.select_token_id(tokens, frame_patch_token_id) + + def get_replacement_audio(item_idx: int) -> PromptUpdateDetails: + """Generate token sequence for an audio.""" + encoder_feats_lengths = out_mm_data.get("encoder_feats_lengths", None) + if encoder_feats_lengths is None: + raise ValueError( + "encoder_feats_lengths missing from processor output; " + "cannot determine audio patch count for prompt replacement." + ) + if isinstance(encoder_feats_lengths, torch.Tensor): + num_patches = int(encoder_feats_lengths[item_idx].item()) + else: + num_patches = encoder_feats_lengths[item_idx] + + # Build token sequence: + tokens: list[int] = [] + tokens.append(audio_start_token_id) + tokens.extend([audio_patch_token_id] * num_patches) + tokens.append(audio_end_token_id) + + # Only tokens receive multimodal embeddings + return PromptUpdateDetails.select_token_id(tokens, audio_patch_token_id) + + # Build prompt updates and process replacement + updates: list[PromptUpdate] = [] + + if "image" in mm_items and mm_items.get_items("image", ImageProcessorItems): + updates.append( + PromptReplacement( + modality="image", + target=PLACEHOLDER_IMAGE_TOKEN_IN_TEXT, + replacement=get_replacement_image, + ) + ) + if "video" in mm_items and mm_items.get_items("video", VideoProcessorItems): + updates.append( + PromptReplacement( + modality="video", + target=PLACEHOLDER_VIDEO_TOKEN_IN_TEXT, + replacement=get_replacement_video, + ) + ) + if "audio" in mm_items and mm_items.get_items("audio", AudioProcessorItems): + updates.append( + PromptReplacement( + modality="audio", + target=PLACEHOLDER_AUDIO_TOKEN_IN_TEXT, + replacement=get_replacement_audio, + ) + ) + return updates + + def _get_mm_fields_config( + self, + hf_inputs: BatchFeature, + hf_processor_mm_kwargs: Mapping[str, object], + ) -> Mapping[str, MultiModalFieldConfig]: + config: dict[str, MultiModalFieldConfig] = {} + + # Image fields, pixel_values is flat (concatenated patches from all images) + image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3))) + if "pixel_values" in hf_inputs: + image_sizes = image_grid_thw.prod(-1) + config["pixel_values"] = MultiModalFieldConfig.flat_from_sizes( + "image", + image_sizes, + ) + if "image_grid_thw" in hf_inputs: + config["image_grid_thw"] = MultiModalFieldConfig.batched("image") + + # Video fields, same flat layout as images + video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3))) + if "pixel_values_videos" in hf_inputs: + video_sizes = video_grid_thw.prod(-1) + config["pixel_values_videos"] = MultiModalFieldConfig.flat_from_sizes( + "video", + video_sizes, + ) + if "video_grid_thw" in hf_inputs: + config["video_grid_thw"] = MultiModalFieldConfig.batched("video") + + # Audio fields + if "audio_feats" in hf_inputs: + config["audio_feats"] = MultiModalFieldConfig.batched("audio") + if "audio_feats_lengths" in hf_inputs: + config["audio_feats_lengths"] = MultiModalFieldConfig.batched("audio") + if "encoder_feats_lengths" in hf_inputs: + config["encoder_feats_lengths"] = MultiModalFieldConfig.batched("audio") + if "placeholder_audio_loc_lens" in hf_inputs: + config["placeholder_audio_loc_lens"] = MultiModalFieldConfig.batched("audio") + + return config + + def _hf_processor_applies_updates( + self, + prompt_text: str, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, object], + tokenization_kwargs: Mapping[str, object], + ) -> bool: + return False + + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + """Call sub-processors for multimodal inputs and tokenize. + + We call the image/audio sub-processors directly (instead of going + through `MingFlashOmniProcessor.__call__`) so that the high-level + placeholder tokens remain **unexpanded** in the tokenized output. + """ + hf_processor = self.info.get_hf_processor() + tokenizer = self.info.get_tokenizer() + + data: dict[str, object] = {} + + images = mm_data.get("images", None) + if images is not None: + image_outputs = hf_processor.image_processor( + images=images, + videos=None, + return_tensors="pt", + ) + data.update(image_outputs) + + videos = mm_data.get("videos", None) + if videos is not None: + video_outputs = hf_processor.image_processor( + images=None, + videos=videos, + return_tensors="pt", + ) + # Rename keys to distinguish from images + if "pixel_values" in video_outputs: + video_outputs["pixel_values_videos"] = video_outputs.pop("pixel_values") + if "image_grid_thw" in video_outputs: + video_outputs["video_grid_thw"] = video_outputs.pop("image_grid_thw") + data.update(video_outputs) + + audios = mm_data.get("audios", None) + if audios is not None: + # vLLM's AudioProcessorItems provides raw numpy arrays (already resampled). + # MingWhisperAudioProcessor expects (waveform, sr) tuples, + # so wrap them with the target sample rate. + target_sr = hf_processor.audio_processor.sampling_rate + audio_tuples = [(a, target_sr) if not isinstance(a, tuple) else a for a in audios] + + audio_outputs = hf_processor.audio_processor( + audio_tuples, + return_tensors="pt", + ) + data.update(audio_outputs) + + # Tokenize text with placeholders still intact + text_outputs = tokenizer(prompt, return_tensors="pt", **tok_kwargs) + data.update(text_outputs) + + return BatchFeature(data=data) + + +@MULTIMODAL_REGISTRY.register_processor( + MingFlashOmniThinkerMultiModalProcessor, + info=MingFlashOmniThinkerProcessingInfo, + dummy_inputs=MingFlashOmniThinkerDummyInputsBuilder, +) +class MingFlashOmniThinkerForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + SupportsMRoPE, + CustomProcessMixin, +): + """Ming Thinker stage: multimodal understanding + (text + image + video + audio) -> text generation. + """ + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={"model.": "language_model."}, + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + # vllm_omni/transformers_utils/processors/ming.py + if modality.startswith("image"): + return "" + elif modality.startswith("video"): + return "