From 0a4e5a8710223a0854f84ad6e7092bfb9aff0c89 Mon Sep 17 00:00:00 2001 From: R2-Y <73573651+R2-Y@users.noreply.github.com> Date: Sun, 30 Nov 2025 04:44:01 +0800 Subject: [PATCH] [Feature] add support for Qwen3-omni (#55) Signed-off-by: Rein Yang Co-authored-by: Rein --- .../stage_configs/qwen2_5_omni.yaml | 2 +- .../offline_inference/qwen2_5_omni/end2end.py | 4 +- .../offline_inference/qwen3_omni/README.md | 34 + .../offline_inference/qwen3_omni/end2end.py | 289 +++ .../qwen3_omni/run_multiple_prompts.sh | 3 + .../qwen3_omni/run_single_prompt.sh | 2 + .../{ => qwen2_5_omni}/README.md | 0 .../{ => qwen2_5_omni}/gradio_demo.py | 0 ...letion_client_for_multimodal_generation.py | 0 .../run_curl_multimodal_generation.sh | 0 examples/online_serving/qwen3_omni/README.md | 77 + .../online_serving/qwen3_omni/gradio_demo.py | 537 ++++++ ...letion_client_for_multimodal_generation.py | 235 +++ .../run_curl_multimodal_generation.sh | 187 ++ vllm_omni/__init__.py | 5 +- vllm_omni/config/model.py | 342 +++- vllm_omni/core/sched/diffusion_scheduler.py | 8 +- vllm_omni/core/sched/generation_scheduler.py | 313 ++++ vllm_omni/engine/__init__.py | 74 +- vllm_omni/engine/arg_utils.py | 36 +- vllm_omni/engine/output_processor.py | 105 +- vllm_omni/entrypoints/async_omni_llm.py | 24 +- vllm_omni/entrypoints/omni_llm.py | 44 +- vllm_omni/entrypoints/omni_stage.py | 115 +- vllm_omni/entrypoints/openai/serving_chat.py | 2 +- vllm_omni/entrypoints/stage_utils.py | 16 +- vllm_omni/model_executor/layers/mrope.py | 16 +- vllm_omni/model_executor/models/__init__.py | 3 + .../model_executor/models/output_templates.py | 13 + .../models/qwen2_5_omni/qwen2_5_omni.py | 13 +- .../models/qwen3_omni/__init__.py | 3 + .../models/qwen3_omni/qwen3_moe.py | 173 ++ .../models/qwen3_omni/qwen3_omni.py | 1132 +++++++++++ .../models/qwen3_omni/qwen3_omni_code2wav.py | 224 +++ .../qwen3_omni_moe_code_predictor_mtp.py | 491 +++++ .../qwen3_omni/qwen3_omni_moe_talker.py | 703 +++++++ .../qwen3_omni/qwen3_omni_moe_thinker.py | 1668 +++++++++++++++++ vllm_omni/model_executor/models/registry.py | 21 + vllm_omni/model_executor/models/utils.py | 9 + .../stage_configs/qwen2_5_omni.yaml | 2 +- .../stage_configs/qwen3_omni_moe.yaml | 94 + .../stage_input_processors/qwen3_omni.py | 171 ++ vllm_omni/patch.py | 12 +- vllm_omni/worker/gpu_ar_model_runner.py | 320 +--- .../worker/gpu_generation_model_runner.py | 582 ++++++ vllm_omni/worker/gpu_generation_worker.py | 84 + vllm_omni/worker/gpu_model_runner.py | 354 +++- 47 files changed, 8064 insertions(+), 478 deletions(-) create mode 100644 examples/offline_inference/qwen3_omni/README.md create mode 100644 examples/offline_inference/qwen3_omni/end2end.py create mode 100644 examples/offline_inference/qwen3_omni/run_multiple_prompts.sh create mode 100644 examples/offline_inference/qwen3_omni/run_single_prompt.sh rename examples/online_serving/{ => qwen2_5_omni}/README.md (100%) rename examples/online_serving/{ => qwen2_5_omni}/gradio_demo.py (100%) rename examples/online_serving/{ => qwen2_5_omni}/openai_chat_completion_client_for_multimodal_generation.py (100%) rename examples/online_serving/{ => qwen2_5_omni}/run_curl_multimodal_generation.sh (100%) create mode 100644 examples/online_serving/qwen3_omni/README.md create mode 100644 examples/online_serving/qwen3_omni/gradio_demo.py create mode 100644 examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py create mode 100644 examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh create mode 100644 vllm_omni/core/sched/generation_scheduler.py create mode 100644 vllm_omni/model_executor/models/output_templates.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/__init__.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py create mode 100644 vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py create mode 100644 vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml create mode 100644 vllm_omni/model_executor/stage_input_processors/qwen3_omni.py create mode 100644 vllm_omni/worker/gpu_generation_model_runner.py create mode 100644 vllm_omni/worker/gpu_generation_worker.py diff --git a/docs/configuration/stage_configs/qwen2_5_omni.yaml b/docs/configuration/stage_configs/qwen2_5_omni.yaml index e4f6b8e7b..c378a5b48 100644 --- a/docs/configuration/stage_configs/qwen2_5_omni.yaml +++ b/docs/configuration/stage_configs/qwen2_5_omni.yaml @@ -30,7 +30,7 @@ stage_args: runtime: process: true devices: "1" - max_batch_size: 3 + max_batch_size: 1 engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index 10bcf4a20..c5be18a94 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -16,7 +16,7 @@ from vllm.multimodal.image import convert_image_mode from vllm.sampling_params import SamplingParams from vllm.utils import FlexibleArgumentParser -from vllm_omni import Omni +from vllm_omni.entrypoints.omni_llm import OmniLLM SEED = 42 @@ -140,7 +140,7 @@ def main(args): model_name = "Qwen/Qwen2.5-Omni-7B" query_result = query_map[args.query_type]() - omni_llm = Omni( + omni_llm = OmniLLM( model=model_name, log_stats=args.enable_stats, log_file=("omni_llm_pipeline.log" if args.enable_stats else None), diff --git a/examples/offline_inference/qwen3_omni/README.md b/examples/offline_inference/qwen3_omni/README.md new file mode 100644 index 000000000..766453667 --- /dev/null +++ b/examples/offline_inference/qwen3_omni/README.md @@ -0,0 +1,34 @@ +# Offline Example of vLLM-Omni for Qwen2.5-omni + +## 🛠️ Installation + +Please refer to [README.md](../../../README.md) + +## Run examples (Qwen3-omni) +### Multiple Prompts +Download dataset from [seed_tts](https://drive.google.com/file/d/1GlSjVfSHkW3-leKKBlfrjuuTGqQ_xaLP/edit). For processing dataset please refer to [Qwen2.5-omni README.md](../qwen2_5_omni/README.md) +Get into the example folder +```bash +cd examples/offline_inference/qwen3_omni +``` +Then run the command below. +```bash +bash run_multiple_prompts.sh +``` +### Single Prompts +Get into the example folder +```bash +cd examples/offline_inference/qwen3_omni +``` +Then run the command below. +```bash +bash run_single_prompt.sh +``` + +### FAQ + +If you encounter error about backend of librosa, try to install ffmpeg with command below. +``` +sudo apt update +sudo apt install ffmpeg +``` diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py new file mode 100644 index 000000000..40b93d549 --- /dev/null +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -0,0 +1,289 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +This example shows how to use vLLM for running offline inference +with the correct prompt format on Qwen3-Omni (thinker only). +""" + +import os +from typing import NamedTuple + +import soundfile as sf +from vllm import SamplingParams +from vllm.assets.audio import AudioAsset +from vllm.assets.image import ImageAsset +from vllm.assets.video import VideoAsset +from vllm.multimodal.image import convert_image_mode +from vllm.utils import FlexibleArgumentParser + +from vllm_omni.entrypoints.omni_llm import OmniLLM + +SEED = 42 + + +class QueryResult(NamedTuple): + inputs: dict + limit_mm_per_prompt: dict[str, int] + + +# NOTE: The default `max_num_seqs` and `max_model_len` may result in OOM on +# lower-end GPUs. +# Unless specified, these settings have been tested to work on a single L4. + +default_system = ( + "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." +) + + +def get_text_query(question: str = None) -> QueryResult: + if question is None: + question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." + prompt = ( + f"<|im_start|>system\n{default_system}<|im_end|>\n" + "<|im_start|>user\n" + f"{question}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + return QueryResult( + inputs={ + "prompt": prompt, + }, + limit_mm_per_prompt={}, + ) + + +def get_video_query(question: str = None) -> QueryResult: + if question is None: + question = "Why is this video funny?" + prompt = ( + f"<|im_start|>system\n{default_system}<|im_end|>\n" + "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>" + f"{question}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": { + "video": VideoAsset(name="baby_reading", num_frames=16).np_ndarrays, + }, + }, + limit_mm_per_prompt={"video": 1}, + ) + + +def get_image_query(question: str = None) -> QueryResult: + if question is None: + question = "What is the content of this image?" + prompt = ( + f"<|im_start|>system\n{default_system}<|im_end|>\n" + "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>" + f"{question}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": { + "image": convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"), + }, + }, + limit_mm_per_prompt={"image": 1}, + ) + + +def get_audio_query(question: str = None) -> QueryResult: + if question is None: + question = "What is the content of this audio?" + prompt = ( + f"<|im_start|>system\n{default_system}<|im_end|>\n" + "<|im_start|>user\n<|audio_start|><|audio_pad|><|audio_end|>" + f"{question}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": { + "audio": AudioAsset("mary_had_lamb").audio_and_sample_rate, + }, + }, + limit_mm_per_prompt={"audio": 1}, + ) + + +query_map = { + "text": get_text_query, + "use_audio": get_audio_query, + "use_image": get_image_query, + "use_video": get_video_query, +} + + +def main(args): + model_name = "/mnt/gh/yrr_vllm/Qwen3-Omni-30B-A3B-Instruct-Full" + query_result = query_map[args.query_type]() + + omni_llm = OmniLLM( + model=model_name, + ) + + thinker_sampling_params = SamplingParams( + temperature=0.4, + top_p=0.9, + top_k=-1, + max_tokens=1200, + repetition_penalty=1.05, + logit_bias={}, + seed=SEED, + ) + + talker_sampling_params = SamplingParams( + temperature=0.9, + top_k=50, + max_tokens=4096, + seed=SEED, + detokenize=False, + repetition_penalty=1.05, + stop_token_ids=[2150], # TALKER_CODEC_EOS_TOKEN_ID + ) + + # Sampling parameters for Code2Wav stage (audio generation) + code2wav_sampling_params = SamplingParams( + temperature=0.0, + top_p=1.0, + top_k=-1, + max_tokens=4096 * 16, + seed=SEED, + detokenize=True, + repetition_penalty=1.1, + ) + + sampling_params_list = [ + thinker_sampling_params, + talker_sampling_params, # code predictor is integrated into talker for Qwen3 Omni + code2wav_sampling_params, + ] + + if args.txt_prompts is None: + prompts = [query_result.inputs for _ in range(args.num_prompts)] + else: + assert args.query_type == "text", "txt-prompts is only supported for text query type" + with open(args.txt_prompts, encoding="utf-8") as f: + lines = [ln.strip() for ln in f.readlines()] + prompts = [get_text_query(ln).inputs for ln in lines if ln != ""] + print(f"[Info] Loaded {len(prompts)} prompts from {args.txt_prompts}") + + omni_outputs = omni_llm.generate(prompts, sampling_params_list) + # Determine output directory: prefer --output-dir; fallback to --output-wav + output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav + os.makedirs(output_dir, exist_ok=True) + + for stage_outputs in omni_outputs: + if stage_outputs.final_output_type == "text": + for output in stage_outputs.request_output: + request_id = int(output.request_id) + text_output = output.outputs[0].text + # Save aligned text file per request + prompt_text = prompts[request_id]["prompt"] + out_txt = os.path.join(output_dir, f"{request_id:05d}.txt") + lines = [] + lines.append("Prompt:\n") + lines.append(str(prompt_text) + "\n") + lines.append("vllm_text_output:\n") + lines.append(str(text_output).strip() + "\n") + try: + with open(out_txt, "w", encoding="utf-8") as f: + f.writelines(lines) + except Exception as e: + print(f"[Warn] Failed writing text file {out_txt}: {e}") + print(f"Request ID: {request_id}, Text saved to {out_txt}") + elif stage_outputs.final_output_type == "audio": + for output in stage_outputs.request_output: + request_id = int(output.request_id) + audio_tensor = output.multimodal_output["audio"] + output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav") + + # Convert to numpy array and ensure correct format + audio_numpy = audio_tensor.float().detach().cpu().numpy() + + # Ensure audio is 1D (flatten if needed) + if audio_numpy.ndim > 1: + audio_numpy = audio_numpy.flatten() + + # Save audio file with explicit WAV format + sf.write(output_wav, audio_numpy, samplerate=24000, format="WAV") + print(f"Request ID: {request_id}, Saved audio to {output_wav}") + + +def parse_args(): + parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser.add_argument( + "--query-type", + "-q", + type=str, + default="mixed_modalities", + choices=query_map.keys(), + help="Query type.", + ) + parser.add_argument( + "--enable-stats", + action="store_true", + default=False, + help="Enable writing detailed statistics (default: disabled)", + ) + parser.add_argument( + "--init-sleep-seconds", + type=int, + default=20, + help="Sleep seconds after starting each stage process to allow initialization (default: 20)", + ) + parser.add_argument( + "--batch-timeout", + type=int, + default=5, + help="Timeout for batching in seconds (default: 5)", + ) + parser.add_argument( + "--init-timeout", + type=int, + default=300, + help="Timeout for initializing stages in seconds (default: 300)", + ) + parser.add_argument( + "--shm-threshold-bytes", + type=int, + default=65536, + help="Threshold for using shared memory in bytes (default: 65536)", + ) + parser.add_argument( + "--output-wav", + default="output_audio", + help="[Deprecated] Output wav directory (use --output-dir).", + ) + parser.add_argument( + "--num-prompts", + type=int, + default=1, + help="Number of prompts to generate.", + ) + parser.add_argument( + "--txt-prompts", + type=str, + default=None, + help="Path to a .txt file with one prompt per line (preferred).", + ) + + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + main(args) + + # use examples: + # python end2end.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct-Full --query-type text + # python end2end.py --model /mnt/gh/yrr_vllm/Qwen3-Omni-30B-A3B-Instruct-Full --query-type use_video_only diff --git a/examples/offline_inference/qwen3_omni/run_multiple_prompts.sh b/examples/offline_inference/qwen3_omni/run_multiple_prompts.sh new file mode 100644 index 000000000..2ec8a1c57 --- /dev/null +++ b/examples/offline_inference/qwen3_omni/run_multiple_prompts.sh @@ -0,0 +1,3 @@ +python end2end.py --output-wav output_audio \ + --query-type text \ + --txt-prompts top10.txt diff --git a/examples/offline_inference/qwen3_omni/run_single_prompt.sh b/examples/offline_inference/qwen3_omni/run_single_prompt.sh new file mode 100644 index 000000000..c6ca09da6 --- /dev/null +++ b/examples/offline_inference/qwen3_omni/run_single_prompt.sh @@ -0,0 +1,2 @@ +python end2end.py --output-wav output_audio \ + --query-type use_audio diff --git a/examples/online_serving/README.md b/examples/online_serving/qwen2_5_omni/README.md similarity index 100% rename from examples/online_serving/README.md rename to examples/online_serving/qwen2_5_omni/README.md diff --git a/examples/online_serving/gradio_demo.py b/examples/online_serving/qwen2_5_omni/gradio_demo.py similarity index 100% rename from examples/online_serving/gradio_demo.py rename to examples/online_serving/qwen2_5_omni/gradio_demo.py diff --git a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py similarity index 100% rename from examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py rename to examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py diff --git a/examples/online_serving/run_curl_multimodal_generation.sh b/examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh similarity index 100% rename from examples/online_serving/run_curl_multimodal_generation.sh rename to examples/online_serving/qwen2_5_omni/run_curl_multimodal_generation.sh diff --git a/examples/online_serving/qwen3_omni/README.md b/examples/online_serving/qwen3_omni/README.md new file mode 100644 index 000000000..01a0fdd28 --- /dev/null +++ b/examples/online_serving/qwen3_omni/README.md @@ -0,0 +1,77 @@ +# Online serving Example of vLLM-omni for Qwen2.5-omni + +## 🛠️ Installation + +Please refer to [README.md](../../README.md) + +## Run examples (Qwen3-Omni) + +Launch the server +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 +``` + +If you have custom stage configs file, launch the server with command below +```bash +vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 --stage-configs-path /path/to/stage_configs_file +``` + +Get into the example folder +```bash +cd examples/online_serving +``` + +Send request via python +```bash +python openai_chat_completion_client_for_multimodal_generation.py --query-type mixed_modalities +``` + +Send request via curl +```bash +bash run_curl_multimodal_generation.sh mixed_modalities +``` + +### FAQ + +If you encounter error about backend of librosa, try to install ffmpeg with command below. +``` +sudo apt update +sudo apt install ffmpeg +``` + +## Run Local Web UI Demo + +This Web UI demo allows users to interact with the model through a web browser. + +### Running Gradio Demo + +Once vllm and vllm-omni are installed, you can launch the web service built on AsyncOmniLLM by + +```bash +python gradio_demo.py --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 7861 +``` + +Then open `http://localhost:7861/` on your local browser to interact with the web UI. + + +### Options + +The gradio demo also supports running with an existing API server and can be customized with the following arguments. + + +```bash +python gradio_demo.py \ + --model Qwen/Qwen3-Omni-30B-A3B-Instruct \ + --use-api-server \ + --api-base http://localhost:8091/v1 \ + --ip 127.0.0.1 \ + --port 7861 +``` + +- `--model`: Model name +- `--use-api-server`: If set, connect to an existing vLLM HTTP API server instead of running AsyncOmniLLM locally. +- `--api-base`: Base URL for vllm serve (only used when `use-api-server` is set, default: http://localhost:8091/v1) +- `--ip`: Host/IP for Gradio server (default: 127.0.0.1) +- `--port`: Port for Gradio server (default: 7861) +- `--stage-configs-path`: Path to custom stage configs YAML file (optional) +- `--share`: Share the Gradio demo publicly (creates a public link) diff --git a/examples/online_serving/qwen3_omni/gradio_demo.py b/examples/online_serving/qwen3_omni/gradio_demo.py new file mode 100644 index 000000000..d6a2237e9 --- /dev/null +++ b/examples/online_serving/qwen3_omni/gradio_demo.py @@ -0,0 +1,537 @@ +import argparse +import os +import random +import signal +import sys +from pathlib import Path +from types import SimpleNamespace +from typing import Any, Optional + +import gradio as gr +import numpy as np +import soundfile as sf +import torch +from PIL import Image +from vllm.assets.video import video_get_metadata, video_to_ndarrays +from vllm.sampling_params import SamplingParams + +from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM + +# Import utils from offline inference example +sys.path.insert(0, os.path.join(os.path.dirname(__file__), "../offline_inference/qwen3_omni")) +try: + from utils import make_omni_prompt +except ImportError: + # Fallback if utils doesn't exist - we'll build prompts directly + make_omni_prompt = None + +SEED = 42 +ASYNC_INIT_TIMEOUT = 600 + +SUPPORTED_MODELS: dict[str, dict[str, Any]] = { + "Qwen/Qwen3-Omni-30B-A3B-Instruct": { + "sampling_params": { + "thinker": { + "temperature": 0.4, + "top_p": 0.9, + "top_k": 1, + "max_tokens": 16384, + "detokenize": True, + "repetition_penalty": 1.05, + "stop_token_ids": [151645], + "seed": SEED, + }, + "talker": { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 4096, + "detokenize": True, + "repetition_penalty": 1.1, + "stop_token_ids": [2150], + }, + "code2wav": { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 2048, + "detokenize": True, + "repetition_penalty": 1.1, + }, + }, + }, +} +# Ensure deterministic behavior across runs. +random.seed(SEED) +np.random.seed(SEED) +torch.manual_seed(SEED) +torch.cuda.manual_seed(SEED) +torch.cuda.manual_seed_all(SEED) +torch.backends.cudnn.deterministic = True +torch.backends.cudnn.benchmark = False +os.environ["PYTHONHASHSEED"] = str(SEED) +os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8" + + +def parse_args(): + parser = argparse.ArgumentParser(description="Gradio demo for Qwen3-Omni online inference.") + parser.add_argument( + "--model", + default="Qwen/Qwen3-Omni-30B-A3B-Instruct", + help="Path to local model checkpoint.", + ) + parser.add_argument( + "--ip", + default="127.0.0.1", + help="Host/IP for gradio `launch`.", + ) + parser.add_argument("--port", type=int, default=7861, help="Port for gradio `launch`.") + parser.add_argument("--share", action="store_true", help="Share the Gradio demo publicly.") + parser.add_argument( + "--stage-configs-path", + type=str, + default=None, + help="Path to custom stage configs YAML file (optional).", + ) + return parser.parse_args() + + +def build_sampling_params(seed: int, model_key: str) -> list[SamplingParams]: + """Build SamplingParams objects by reusing the dict definitions.""" + return [SamplingParams(**params_dict) for params_dict in build_sampling_params_dict(seed, model_key)] + + +def build_sampling_params_dict(seed: int, model_key: str) -> list[dict]: + """Build sampling params as dict for HTTP API mode.""" + model_conf = SUPPORTED_MODELS.get(model_key) + if model_conf is None: + raise ValueError(f"Unsupported model '{model_key}'") + + sampling_templates: dict[str, dict[str, Any]] = model_conf["sampling_params"] + sampling_params: list[dict] = [] + for stage_name, template in sampling_templates.items(): + params = dict(template) + params["seed"] = seed + sampling_params.append(params) + return sampling_params + + +def create_prompt_args(base_args: argparse.Namespace) -> SimpleNamespace: + # The prompt builder expects a minimal namespace with these attributes. + return SimpleNamespace( + model=base_args.model, + prompt_type="text", + tokenize=True, + use_torchvision=True, + legacy_omni_video=False, + ) + + +def process_audio_file( + audio_file: Optional[Any], +) -> Optional[tuple[np.ndarray, int]]: + """Normalize Gradio audio input to (np.ndarray, sample_rate).""" + if audio_file is None: + return None + + sample_rate: Optional[int] = None + audio_np: Optional[np.ndarray] = None + + def _load_from_path(path_str: str) -> Optional[tuple[np.ndarray, int]]: + if not path_str: + return None + path = Path(path_str) + if not path.exists(): + return None + data, sr = sf.read(path) + if data.ndim > 1: + data = data[:, 0] + return data.astype(np.float32), int(sr) + + if isinstance(audio_file, tuple): + if len(audio_file) == 2: + first, second = audio_file + # Case 1: (sample_rate, np.ndarray) + if isinstance(first, (int, float)) and isinstance(second, np.ndarray): + sample_rate = int(first) + audio_np = second + # Case 2: (filepath, (sample_rate, np.ndarray or list)) + elif isinstance(first, str): + if isinstance(second, tuple) and len(second) == 2: + sr_candidate, data_candidate = second + if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray): + sample_rate = int(sr_candidate) + audio_np = data_candidate + if audio_np is None: + loaded = _load_from_path(first) + if loaded is not None: + audio_np, sample_rate = loaded + # Case 3: (None, (sample_rate, np.ndarray)) + elif first is None and isinstance(second, tuple) and len(second) == 2: + sr_candidate, data_candidate = second + if isinstance(sr_candidate, (int, float)) and isinstance(data_candidate, np.ndarray): + sample_rate = int(sr_candidate) + audio_np = data_candidate + elif len(audio_file) == 1 and isinstance(audio_file[0], str): + loaded = _load_from_path(audio_file[0]) + if loaded is not None: + audio_np, sample_rate = loaded + elif isinstance(audio_file, str): + loaded = _load_from_path(audio_file) + if loaded is not None: + audio_np, sample_rate = loaded + + if audio_np is None or sample_rate is None: + return None + + if audio_np.ndim > 1: + audio_np = audio_np[:, 0] + + return audio_np.astype(np.float32), sample_rate + + +def process_image_file(image_file: Optional[Image.Image]) -> Optional[Image.Image]: + """Process image file from Gradio input. + + Returns: + PIL Image in RGB mode or None if no image provided. + """ + if image_file is None: + return None + # Convert to RGB if needed + if image_file.mode != "RGB": + image_file = image_file.convert("RGB") + return image_file + + +def process_video_file( + video_file: Optional[str], + enable_audio_in_video: bool = False, + max_frames: int = 32, +) -> Optional[tuple[np.ndarray, dict[str, Any], Optional[tuple[np.ndarray, int]]]]: + """Process video file and optionally extract audio track.""" + if video_file is None: + return None + + video_path = Path(video_file) + if not video_path.exists(): + print(f"Video file not found: {video_path}") + return None + + try: + frames = video_to_ndarrays(str(video_path), num_frames=max_frames) + metadata = video_get_metadata(str(video_path), num_frames=frames.shape[0]) + except Exception as exc: # pylint: disable=broad-except + print(f"Failed to decode video {video_path}: {exc}") + return None + + audio_tuple: Optional[tuple[np.ndarray, int]] = None + if enable_audio_in_video: + try: + import librosa # type: ignore import + + audio_signal, sampling_rate = librosa.load(str(video_path), sr=16000) + audio_tuple = (audio_signal.astype(np.float32), sampling_rate) + except Exception as exc: # pylint: disable=broad-except + print(f"Failed to extract audio from video {video_path}: {exc}") + + return frames, metadata, audio_tuple + + +async def run_inference_async_omni_llm( + omni_llm: AsyncOmniLLM, + sampling_params: list[SamplingParams], + prompt_args_template: SimpleNamespace, + user_prompt: str, + audio_file: Optional[tuple[str, tuple[int, np.ndarray]]] = None, + image_file: Optional[Image.Image] = None, + video_file: Optional[str] = None, + use_audio_in_video: bool = False, +): + """Run inference using AsyncOmniLLM directly with multimodal support.""" + if not user_prompt.strip() and not audio_file and not image_file and not video_file: + return "Please provide at least a text prompt or multimodal input.", None + + try: + # Build prompt with multimodal data + prompt_args = SimpleNamespace(**prompt_args_template.__dict__) + + # Process multimodal inputs + multi_modal_data = {} + mm_processor_kwargs = {} + + # Process audio + audio_data = process_audio_file(audio_file) + if audio_data is not None: + multi_modal_data["audio"] = audio_data + + # Process image + image_data = process_image_file(image_file) + if image_data is not None: + multi_modal_data["image"] = image_data + + # Process video + video_payload = process_video_file(video_file, enable_audio_in_video=use_audio_in_video) + if video_payload is not None: + video_frames, video_metadata, extracted_audio = video_payload + video_entry: Any + if video_metadata: + video_entry = (video_frames, video_metadata) + else: + video_entry = video_frames + multi_modal_data["video"] = video_entry + if use_audio_in_video and extracted_audio is not None and "audio" not in multi_modal_data: + multi_modal_data["audio"] = extracted_audio + mm_processor_kwargs["use_audio_in_video"] = True + + # Build the prompt input + if make_omni_prompt is not None: + omni_prompt = make_omni_prompt(prompt_args, user_prompt) + # Add multimodal data if present + if multi_modal_data: + if isinstance(omni_prompt, dict): + omni_prompt["multi_modal_data"] = multi_modal_data + if mm_processor_kwargs: + omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs + else: + # If make_omni_prompt returns a string, we need to wrap it + omni_prompt = { + "prompt": omni_prompt, + "multi_modal_data": multi_modal_data, + } + if mm_processor_kwargs: + omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs + else: + # Fallback: build prompt directly + default_system = ( + "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." + ) + prompt = f"<|im_start|>system\n{default_system}<|im_end|>\n<|im_start|>user\n" + if audio_data: + prompt += "<|audio_bos|><|AUDIO|><|audio_eos|>" + if image_data: + prompt += "<|vision_bos|><|IMAGE|><|vision_eos|>" + if video_payload is not None: + prompt += "<|vision_bos|><|VIDEO|><|vision_eos|>" + if user_prompt.strip(): + prompt += f"{user_prompt}" + prompt += "<|im_end|>\n<|im_start|>assistant\n" + + omni_prompt = { + "prompt": prompt, + "multi_modal_data": multi_modal_data, + } + if mm_processor_kwargs: + omni_prompt["mm_processor_kwargs"] = mm_processor_kwargs + + request_id = "0" + text_outputs: list[str] = [] + audio_output = None + + async for stage_outputs in omni_llm.generate( + prompt=omni_prompt, + request_id=request_id, + sampling_params_list=sampling_params, + ): + # stage_outputs.request_output is a RequestOutput object, not a list + request_output = stage_outputs.request_output + if stage_outputs.final_output_type == "text": + if request_output.outputs: + for output in request_output.outputs: + if output.text: + text_outputs.append(output.text) + elif stage_outputs.final_output_type == "audio": + # multimodal_output is on the RequestOutput object + # See vllm_omni/entrypoints/openai/serving_chat.py:680 for reference + if hasattr(request_output, "multimodal_output") and request_output.multimodal_output: + if "audio" in request_output.multimodal_output: + audio_tensor = request_output.multimodal_output["audio"] + # Ensure audio is 1D (flatten if needed) + if hasattr(audio_tensor, "ndim") and audio_tensor.ndim > 1: + audio_tensor = audio_tensor.flatten() + audio_np = audio_tensor.detach().cpu().numpy() + audio_output = ( + 24000, # sampling rate in Hz + audio_np, + ) + + text_response = "\n\n".join(text_outputs) if text_outputs else "No text output." + return text_response, audio_output + except Exception as exc: # pylint: disable=broad-except + return f"Inference failed: {exc}", None + + +def build_interface( + omni_llm: AsyncOmniLLM, + sampling_params: list[SamplingParams], + prompt_args_template: SimpleNamespace, + model: str, +): + """Build Gradio interface for AsyncOmniLLM mode.""" + + async def run_inference( + user_prompt: str, + audio_file: Optional[tuple[str, tuple[int, np.ndarray]]], + image_file: Optional[Image.Image], + video_file: Optional[str], + use_audio_in_video: bool, + ): + return await run_inference_async_omni_llm( + omni_llm, + sampling_params, + prompt_args_template, + user_prompt, + audio_file, + image_file, + video_file, + use_audio_in_video, + ) + + css = """ + .media-input-container { + display: flex; + gap: 10px; + } + .media-input-container > div { + flex: 1; + } + .media-input-container .image-input, + .media-input-container .audio-input { + height: 300px; + } + .media-input-container .video-column { + height: 300px; + display: flex; + flex-direction: column; + } + .media-input-container .video-input { + flex: 1; + min-height: 0; + } + #generate-btn button { + width: 100%; + } + """ + + with gr.Blocks(css=css) as demo: + gr.Markdown("# vLLM-Omni Online Serving Demo") + gr.Markdown(f"**Model:** {model} \n\n") + + with gr.Column(): + with gr.Row(): + input_box = gr.Textbox( + label="Text Prompt", + placeholder="For example: Describe what happens in the media inputs.", + lines=4, + scale=1, + ) + with gr.Row(elem_classes="media-input-container"): + image_input = gr.Image( + label="Image Input (optional)", + type="pil", + sources=["upload"], + scale=1, + elem_classes="image-input", + ) + with gr.Column(scale=1, elem_classes="video-column"): + video_input = gr.Video( + label="Video Input (optional)", + sources=["upload"], + elem_classes="video-input", + ) + use_audio_in_video_checkbox = gr.Checkbox( + label="Use audio from video", + value=False, + info="Extract the video's audio track when provided.", + ) + audio_input = gr.Audio( + label="Audio Input (optional)", + type="numpy", + sources=["upload", "microphone"], + scale=1, + elem_classes="audio-input", + ) + + with gr.Row(): + generate_btn = gr.Button( + "Generate", + variant="primary", + size="lg", + elem_id="generate-btn", + ) + + with gr.Row(): + text_output = gr.Textbox(label="Text Output", lines=10, scale=2) + audio_output = gr.Audio(label="Audio Output", interactive=False, scale=1) + + generate_btn.click( + fn=run_inference, + inputs=[input_box, audio_input, image_input, video_input, use_audio_in_video_checkbox], + outputs=[text_output, audio_output], + ) + demo.queue() + return demo + + +def main(): + args = parse_args() + omni_llm = None + + model_name = "/".join(args.model.split("/")[-2:]) + assert model_name in SUPPORTED_MODELS, ( + f"Unsupported model '{model_name}'. Supported models: {SUPPORTED_MODELS.keys()}" + ) + + # Register signal handlers for graceful shutdown + def signal_handler(sig, frame): + print("\nReceived interrupt signal, shutting down...") + if omni_llm is not None: + try: + omni_llm.shutdown() + except Exception as e: + print(f"Error during shutdown: {e}") + sys.exit(0) + + signal.signal(signal.SIGINT, signal_handler) + signal.signal(signal.SIGTERM, signal_handler) + + print(f"Initializing AsyncOmniLLM with model: {args.model}") + if args.stage_configs_path: + print(f"Using custom stage configs: {args.stage_configs_path}") + + sampling_params = build_sampling_params(SEED, model_name) + omni_llm = AsyncOmniLLM( + model=args.model, + stage_configs_path=args.stage_configs_path, + init_timeout=ASYNC_INIT_TIMEOUT, + ) + print("✓ AsyncOmniLLM initialized successfully") + prompt_args_template = create_prompt_args(args) + + demo = build_interface( + omni_llm, + sampling_params, + prompt_args_template, + args.model, + ) + try: + demo.launch( + server_name=args.ip, + server_port=args.port, + share=args.share, + ) + except KeyboardInterrupt: + print("\nShutting down...") + finally: + # Cleanup + if omni_llm is not None: + try: + omni_llm.shutdown() + except Exception as e: + print(f"Error during cleanup: {e}") + + +if __name__ == "__main__": + main() diff --git a/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py new file mode 100644 index 000000000..d4522b2a4 --- /dev/null +++ b/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py @@ -0,0 +1,235 @@ +import base64 +from typing import NamedTuple + +import requests +from openai import OpenAI +from vllm.assets.audio import AudioAsset +from vllm.utils import FlexibleArgumentParser + +# Modify OpenAI's API key and API base to use vLLM's API server. +openai_api_key = "EMPTY" +openai_api_base = "http://localhost:8091/v1" + +client = OpenAI( + # defaults to os.environ.get("OPENAI_API_KEY") + api_key=openai_api_key, + base_url=openai_api_base, +) + +SEED = 42 + + +class QueryResult(NamedTuple): + inputs: dict + limit_mm_per_prompt: dict[str, int] + + +def encode_base64_content_from_url(content_url: str) -> str: + """Encode a content retrieved from a remote url to base64 format.""" + + with requests.get(content_url) as response: + response.raise_for_status() + result = base64.b64encode(response.content).decode("utf-8") + + return result + + +def get_system_prompt(): + return { + "role": "system", + "content": [ + { + "type": "text", + "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." + ), + } + ], + } + + +def get_text_query(): + question = "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." + prompt = { + "role": "user", + "content": [ + { + "type": "text", + "text": f"{question}", + } + ], + } + return prompt + + +default_system = ( + "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." +) + + +def get_video_query() -> QueryResult: + question = "Why is this video funny?" + prompt = { + "role": "user", + "content": [ + { + "type": "video_url", + "video_url": { + "url": "https://huggingface.co/datasets/raushan-testing-hf/videos-test/resolve/main/sample_demo_1.mp4" + }, + }, + { + "type": "text", + "text": f"{question}", + }, + ], + } + return prompt + + +def get_image_query() -> QueryResult: + question = "What is the content of this image?" + prompt = { + "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": f"{question}", + }, + ], + } + return prompt + + +def get_audio_query() -> QueryResult: + question = "What is the content of this audio?" + prompt = { + "role": "user", + "content": [ + { + "type": "audio_url", + "audio_url": {"url": AudioAsset("mary_had_lamb").url}, + }, + { + "type": "text", + "text": f"{question}", + }, + ], + } + return prompt + + +query_map = { + "text": get_text_query, + "use_audio": get_audio_query, + "use_image": get_image_query, + "use_video": get_video_query, +} + + +def run_multimodal_generation(args) -> None: + model_name = args.model + thinker_sampling_params = { + "temperature": 0.4, # Deterministic + "top_p": 0.9, + "top_k": 1, + "max_tokens": 16384, + "repetition_penalty": 1.05, + "stop_token_ids": [151645], # Qwen EOS token <|im_end|> + "seed": SEED, + } + + # Sampling parameters for Talker stage (codec generation) + # Stop at codec EOS token + talker_sampling_params = { + "temperature": 0.9, + "top_k": 50, + "max_tokens": 4096, + "seed": SEED, + "detokenize": False, + "repetition_penalty": 1.05, + "stop_token_ids": [2150], # TALKER_CODEC_EOS_TOKEN_ID + } + + # # Sampling parameters for Code2Wav stage (audio generation) + code2wav_sampling_params = { + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 4096 * 16, + "seed": SEED, + "detokenize": True, + "repetition_penalty": 1.1, + } + + sampling_params_list = [ + thinker_sampling_params, + talker_sampling_params, + code2wav_sampling_params, + ] + + prompt = query_map[args.query_type]() + extra_body = { + "sampling_params_list": sampling_params_list # Optional, it has a default setting in stage_configs of the corresponding model. + } + + if args.query_type == "use_audio_in_video": + extra_body["mm_processor_kwargs"] = {"use_audio_in_video": True} + + chat_completion = client.chat.completions.create( + messages=[ + get_system_prompt(), + prompt, + ], + model=model_name, + extra_body=extra_body, + ) + + count = 0 + for choice in chat_completion.choices: + print(f"choice: {choice}") + if choice.message.audio: + audio_data = base64.b64decode(choice.message.audio.data) + audio_file_path = f"audio_{count}.wav" + with open(audio_file_path, "wb") as f: + f.write(audio_data) + print(f"Audio saved to {audio_file_path}") + count += 1 + elif choice.message.content: + print("Chat completion output from text:", choice.message.content) + + +def parse_args(): + parser = FlexibleArgumentParser(description="Demo on using vLLM for offline inference with audio language models") + parser.add_argument( + "--query-type", + "-q", + type=str, + default="use_video", + choices=query_map.keys(), + help="Query type.", + ) + parser.add_argument( + "--model", + "-m", + type=str, + default="Qwen/Qwen3-Omni-30B-A3B-Instruct", + help="Model Name / Path", + ) + + return parser.parse_args() + + +if __name__ == "__main__": + args = parse_args() + run_multimodal_generation(args) diff --git a/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh b/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh new file mode 100644 index 000000000..95c85370c --- /dev/null +++ b/examples/online_serving/qwen3_omni/run_curl_multimodal_generation.sh @@ -0,0 +1,187 @@ +#!/usr/bin/env bash +set -euo pipefail + +# Default query type +QUERY_TYPE="${1:-mixed_modalities}" + +# Validate query type +if [[ ! "$QUERY_TYPE" =~ ^(mixed_modalities|use_audio_in_video|multi_audios|text)$ ]]; then + echo "Error: Invalid query type '$QUERY_TYPE'" + echo "Usage: $0 [mixed_modalities|use_audio_in_video|multi_audios|text]" + echo " mixed_modalities: Audio + Image + Video + Text query" + echo " use_audio_in_video: Video + Text query (with audio extraction from video)" + echo " multi_audios: Two audio clips + Text query" + echo " text: Text query" + exit 1 +fi + +SEED=42 + +thinker_sampling_params='{ + "temperature": 0.4, + "top_p": 0.9, + "top_k": 1, + "max_tokens": 16384, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.05, + "stop_token_ids": [151645] +}' + +talker_sampling_params='{ + "temperature": 0.9, + "top_k": 50, + "max_tokens": 4096, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.05, + "stop_token_ids": [2150] +}' + +code2wav_sampling_params='{ + "temperature": 0.0, + "top_p": 1.0, + "top_k": -1, + "max_tokens": 4096 * 16, + "seed": 42, + "detokenize": true, + "repetition_penalty": 1.1 +}' +# 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" +WINNING_CALL_AUDIO_URL="https://vllm-public-assets.s3.us-west-2.amazonaws.com/multimodal_asset/winning_call.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 and extra fields based on query type +case "$QUERY_TYPE" in + text) + user_content='[ + { + "type": "text", + "text": "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." + } + ]' + sampling_params_list='[ + '"$thinker_sampling_params"', + '"$talker_sampling_params"', + '"$code2wav_sampling_params"' + ]' + mm_processor_kwargs="{}" + ;; + mixed_modalities) + user_content='[ + { + "type": "audio_url", + "audio_url": { + "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'" + } + }, + { + "type": "image_url", + "image_url": { + "url": "'"$CHERRY_BLOSSOM_IMAGE_URL"'" + } + }, + { + "type": "video_url", + "video_url": { + "url": "'"$SAMPLE_VIDEO_URL"'" + } + }, + { + "type": "text", + "text": "What is recited in the audio? What is the content of this image? Why is this video funny?" + } + ]' + sampling_params_list='[ + '"$thinker_sampling_params"', + '"$talker_sampling_params"', + '"$code2wav_sampling_params"' + ]' + mm_processor_kwargs="{}" + ;; + use_audio_in_video) + user_content='[ + { + "type": "video_url", + "video_url": { + "url": "'"$SAMPLE_VIDEO_URL"'" + } + }, + { + "type": "text", + "text": "Describe the content of the video, then convert what the baby say into text." + } + ]' + sampling_params_list='[ + '"$thinker_sampling_params"', + '"$talker_sampling_params"', + '"$code2wav_sampling_params"' + ]' + mm_processor_kwargs='{ + "use_audio_in_video": true + }' + ;; + multi_audios) + user_content='[ + { + "type": "audio_url", + "audio_url": { + "url": "'"$MARY_HAD_LAMB_AUDIO_URL"'" + } + }, + { + "type": "audio_url", + "audio_url": { + "url": "'"$WINNING_CALL_AUDIO_URL"'" + } + }, + { + "type": "text", + "text": "Are these two audio clips the same?" + } + ]' + sampling_params_list='[ + '"$thinker_sampling_params"', + '"$talker_sampling_params"', + '"$code2wav_sampling_params"' + ]' + mm_processor_kwargs="{}" + ;; +esac + +echo "Running query type: $QUERY_TYPE" +echo "" + + +output=$(curl -sS -X POST http://localhost:8091/v1/chat/completions \ + -H "Content-Type: application/json" \ + -d @- < list[str]: return [self.model_arch] + + def draw_hf_text_config(self): + # transformers' get_text_config method is used to get the text config from thinker_config. + # to handle the case that each model stage has their own text config, + # we need to draw the text config from the corresponding model stage. + if self.hf_config_name is None: + return get_hf_text_config(self.hf_config) + return getattr(self.hf_config, self.hf_config_name).get_text_config() + + def __post_init__( + self, + # Multimodal config init vars + limit_mm_per_prompt: Optional[dict[str, int]], + media_io_kwargs: Optional[dict[str, dict[str, Any]]], + mm_processor_kwargs: Optional[dict[str, Any]], + mm_processor_cache_gb: Optional[float], + mm_processor_cache_type: Optional[MMCacheType], + mm_shm_cache_max_object_size_mb: Optional[int], + mm_encoder_tp_mode: Optional[MMEncoderTPMode], + interleave_mm_strings: Optional[bool], + skip_mm_profiling: Optional[bool], + video_pruning_rate: Optional[float], + ) -> None: + # Set the default seed to 0 in V1. + if envs.VLLM_USE_V1 and self.seed is None: + self.seed = 0 + if not envs.VLLM_ENABLE_V1_MULTIPROCESSING: + logger.warning( + "The global random seed is set to %d. Since " + "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may " + "affect the random state of the Python process that " + "launched vLLM.", + self.seed, + ) + + # Keep set served_model_name before maybe_model_redirect(self.model) + self.served_model_name = get_served_model_name(self.model, self.served_model_name) + self.model = maybe_model_redirect(self.model) + # The tokenizer is consistent with the model by default. + if self.tokenizer is None: + self.tokenizer = self.model + if self.tokenizer_revision is None: + self.tokenizer_revision = self.revision + self.tokenizer = maybe_model_redirect(self.tokenizer) + + if isinstance(self.hf_config_path, str): + self.hf_config_path = maybe_model_redirect(self.hf_config_path) + + if callable(self.hf_overrides): + hf_overrides_kw: dict[str, Any] = {} + hf_overrides_fn = self.hf_overrides + dict_overrides: dict[str, Any] = {} + else: + # Separate dict overrides from flat ones; convert OmegaConf to dict + hf_overrides_kw = {} + dict_overrides = {} + for key, value in self.hf_overrides.items(): + if isinstance(value, DictConfig): + value = OmegaConf.to_container(value, resolve=True, structured_config_mode=None) + if isinstance(value, dict): + dict_overrides[key] = value + else: + hf_overrides_kw[key] = value + hf_overrides_fn = None + + if self.rope_scaling: + hf_override: dict[str, Any] = {"rope_scaling": self.rope_scaling} + hf_overrides_kw.update(hf_override) + hf_overrides_str = json.dumps(hf_overrides_kw) + msg = ( + "`--rope-scaling` will be removed in a future release. " + f"'Please instead use `--hf-overrides '{hf_overrides_str}'`" + ) + warnings.warn(DeprecationWarning(msg), stacklevel=2) + if self.rope_theta is not None: + hf_override = {"rope_theta": self.rope_theta} + hf_overrides_kw.update(hf_override) + hf_overrides_str = json.dumps(hf_overrides_kw) + msg = ( + "`--rope-theta` will be removed in a future release. " + f"'Please instead use `--hf-overrides '{hf_overrides_str}'`" + ) + warnings.warn(DeprecationWarning(msg), stacklevel=2) + + self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer) + + if (backend := envs.VLLM_ATTENTION_BACKEND) and backend == "FLASHINFER" and find_spec("flashinfer") is None: + raise ValueError( + "VLLM_ATTENTION_BACKEND is set to FLASHINFER, but flashinfer " + "module was not found. See " + "https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile " + "for instructions on how to install it." + ) + + if self.override_attention_dtype is not None and not current_platform.is_rocm(): + warnings.warn("override-attention-dtype is set but not using ROCm platform", stacklevel=2) + + if self.enable_sleep_mode and not current_platform.is_sleep_mode_available(): + raise ValueError("Sleep mode is not supported on current platform.") + + hf_config = get_config( + self.hf_config_path or self.model, + self.trust_remote_code, + self.revision, + self.code_revision, + self.config_format, + hf_overrides_kw=hf_overrides_kw, + hf_overrides_fn=hf_overrides_fn, + ) + + if dict_overrides: + self._apply_dict_overrides(hf_config, dict_overrides) + + self.hf_config = hf_config + self.hf_text_config = self.draw_hf_text_config() + self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None) + self.encoder_config = self._get_encoder_config() + self.hf_image_processor_config = get_hf_image_processor_config( + self.model, hf_token=self.hf_token, revision=self.revision + ) + + architectures = self.architectures + registry = self.registry + is_generative_model = registry.is_text_generation_model(architectures, self) + is_pooling_model = registry.is_pooling_model(architectures, self) + + def _task_to_convert(task: TaskOption) -> ConvertType: + if task == "embedding" or task == "embed": + return "embed" + if task == "classify": + return "classify" + if task == "reward": + return "reward" + if task == "score": + new_task = self._get_default_pooling_task(architectures) + return "classify" if new_task == "classify" else "embed" + + return "none" + + if self.task is not None: + runner: RunnerOption = "auto" + convert: ConvertType | Literal["auto"] = "auto" # type: ignore[name-defined] + msg_prefix = ( + "The 'task' option has been deprecated and will be removed in v0.13.0 or v1.0, whichever comes first." + ) + msg_hint = "Please remove this option." + + is_generative_task = self.task in _RUNNER_TASKS["generate"] + is_pooling_task = self.task in _RUNNER_TASKS["pooling"] + + if is_generative_model and is_pooling_model: + if is_generative_task: + runner = "generate" + convert = "auto" # type: ignore[assignment] + msg_hint = ( + "Please replace this option with `--runner " + "generate` to continue using this model " + "as a generative model." + ) + elif is_pooling_task: + runner = "pooling" + convert = "auto" # type: ignore[assignment] + msg_hint = ( + "Please replace this option with `--runner " + "pooling` to continue using this model " + "as a pooling model." + ) + else: # task == "auto" + pass + elif is_generative_model or is_pooling_model: + if is_generative_task: + runner = "generate" + convert = "auto" # type: ignore[assignment] + msg_hint = "Please remove this option" + elif is_pooling_task: + runner = "pooling" + convert = _task_to_convert(self.task) + msg_hint = ( + "Please replace this option with `--convert " + f"{convert}` to continue using this model " + "as a pooling model." + ) + else: # task == "auto" + pass + else: + raise AssertionError( + f"The model should be a generative or pooling model when task is set to {self.task!r}." + ) + + self.runner = runner + self.convert = convert # type: ignore[assignment] + + msg = f"{msg_prefix} {msg_hint}" + warnings.warn(msg, DeprecationWarning, stacklevel=2) + + self.runner_type = self._get_runner_type(architectures, self.runner) + self.convert_type = self._get_convert_type(architectures, self.runner_type, self.convert) + + if self.runner_type == "generate" and not is_generative_model: + generate_converts = _RUNNER_CONVERTS["generate"] + if self.convert_type not in generate_converts: + raise ValueError("This model does not support `--runner generate`.") + if self.runner_type == "pooling" and not is_pooling_model: + pooling_converts = _RUNNER_CONVERTS["pooling"] + if self.convert_type not in pooling_converts: + convert_option = "<" + "|".join(pooling_converts) + ">" + raise ValueError( + "This model does not support `--runner pooling`. " + f"You can pass `--convert {convert_option} to adapt " + "it into a pooling model." + ) + + # Note: Initialize these attributes early because transformers fallback + # may fail to load dynamic modules in child processes + model_info, arch = registry.inspect_model_cls(architectures, self) + self._model_info = model_info + self._architecture = arch + logger.info("Resolved architecture: %s", arch) + + # Init pooler config if needed + if self.runner_type == "pooling": + if self.override_pooler_config is not None: + logger.warning_once( + "`override_pooler_config` is deprecated and will be " + "removed in v0.12.0 or v1.0.0, whichever is sooner. " + "Please use `pooler_config` instead." + ) + + if isinstance(self.override_pooler_config, dict): + self.pooler_config = PoolerConfig(**self.override_pooler_config) + else: + self.pooler_config = self.override_pooler_config + + if self.pooler_config is None: + self.pooler_config = PoolerConfig() + + base_config = get_pooling_config(self.model, self.revision) + if base_config is not None: + for k, v in base_config.items(): + if getattr(self.pooler_config, k) is None: + setattr(self.pooler_config, k, v) + + default_pooling_type = self._model_info.default_pooling_type + if self.pooler_config.pooling_type is None: + self.pooler_config.pooling_type = default_pooling_type + + self.dtype: torch.dtype = _get_and_verify_dtype( + self.model, + self.hf_config, + self.dtype, + is_pooling_model=self.runner_type == "pooling", + revision=self.revision, + ) + + # Interleaved attention is not supported by some backends in V0 + if ( + not self.disable_sliding_window + and is_interleaved(self.hf_text_config) + and not envs.VLLM_USE_V1 + and (backend := envs.VLLM_ATTENTION_BACKEND) in ("XFORMERS", "FLASHINFER") + ): + logger.warning_once( + "%s has interleaved attention, which is currently not " + "supported by the %s backend. Disabling sliding window and " + "capping the max length to the sliding window size (%d).", + self.hf_text_config.model_type, + backend, + self.hf_text_config.sliding_window, + ) + self.disable_sliding_window = True + + self.original_max_model_len = self.max_model_len + self.max_model_len = self.get_and_verify_max_len(self.max_model_len) + # Init multimodal config if needed + if self._model_info.supports_multimodal: + if mm_encoder_tp_mode == "data" and not self._model_info.supports_multimodal_encoder_tp_data: + logger.warning_once( + "This model does not support `--mm-encoder-tp-mode data`. " + "Falling back to `--mm-encoder-tp-mode weights`." + ) + mm_encoder_tp_mode = "weights" + + mm_config_kwargs = dict( + limit_per_prompt=limit_mm_per_prompt, + media_io_kwargs=media_io_kwargs, + mm_processor_kwargs=mm_processor_kwargs, + mm_processor_cache_gb=mm_processor_cache_gb, + mm_processor_cache_type=mm_processor_cache_type, + mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb, + mm_encoder_tp_mode=mm_encoder_tp_mode, + interleave_mm_strings=interleave_mm_strings, + skip_mm_profiling=skip_mm_profiling, + video_pruning_rate=video_pruning_rate, + ) + + mm_config_kwargs = {k: v for k, v in mm_config_kwargs.items() if v is not None} + + self.multimodal_config = MultiModalConfig(**mm_config_kwargs) + + if self.disable_sliding_window: + self.hf_text_config.sliding_window = None + + if not self.skip_tokenizer_init: + self._verify_tokenizer_mode() + + self.config_updated = False + + self._verify_quantization() + self._verify_cuda_graph() + self._verify_bnb_config() diff --git a/vllm_omni/core/sched/diffusion_scheduler.py b/vllm_omni/core/sched/diffusion_scheduler.py index 1b960f937..e8707aabf 100644 --- a/vllm_omni/core/sched/diffusion_scheduler.py +++ b/vllm_omni/core/sched/diffusion_scheduler.py @@ -249,6 +249,12 @@ class DiffusionScheduler(OmniScheduler): # Get prompt logprobs for this request. prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) + # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput + pooling_output_dict = None + if pooler_output is not None: + # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type + # which expects Optional[dict[str, torch.Tensor]] + pooling_output_dict = {"model_outputs": pooler_output} if new_token_ids or pooler_output is not None or kv_transfer_params: # Add EngineCoreOutput for this Request. outputs[request.client_index].append( @@ -258,7 +264,7 @@ class DiffusionScheduler(OmniScheduler): finish_reason=request.get_finished_reason(), new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, - pooling_output=pooler_output, + pooling_output=pooling_output_dict, stop_reason=request.stop_reason, events=request.take_events(), kv_transfer_params=kv_transfer_params, diff --git a/vllm_omni/core/sched/generation_scheduler.py b/vllm_omni/core/sched/generation_scheduler.py new file mode 100644 index 000000000..7916c36ed --- /dev/null +++ b/vllm_omni/core/sched/generation_scheduler.py @@ -0,0 +1,313 @@ +import time +from collections import defaultdict +from typing import Optional + +from vllm.v1.core.kv_cache_manager import KVCacheBlocks +from vllm.v1.core.sched.request_queue import create_request_queue +from vllm.v1.core.sched.scheduler import ( + EngineCoreOutputs, + Request, + RequestStatus, + SchedulerOutput, + SpecDecodingStats, +) +from vllm.v1.core.sched.utils import remove_all +from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput + +from vllm_omni.core.sched.output import OmniNewRequestData +from vllm_omni.core.sched.scheduler import OmniScheduler +from vllm_omni.outputs import OmniModelRunnerOutput + + +class GenerationScheduler(OmniScheduler): + def schedule(self) -> SchedulerOutput: + """Diffusion fast path: + - Feed all input tokens of the request at once + (if 0, allocate 1 placeholder token). + - If the token budget cannot be satisfied at once, fall back to the + default vLLM scheduling. + """ + + # Select requests with zero prompt and using pooling + # (diffusion results returned via pooler_output) + token_budget = self.max_num_scheduled_tokens + capacity = self.max_num_running_reqs - len(self.running) + scheduled_timestamp = time.monotonic() + + scheduled_new_reqs: list[Request] = [] + scheduled_resumed_reqs: list[Request] = [] + scheduled_running_reqs: list[Request] = [] + + req_to_new_blocks: dict[str, KVCacheBlocks] = {} + num_scheduled_tokens: dict[str, int] = {} + scheduled_spec_decode_tokens: dict[str, list[int]] = {} + scheduled_encoder_inputs: dict[str, list[int]] = {} + + # Temporary queue: preserve waiting order, do not disturb non-diffusion requests + skipped_waiting_requests = create_request_queue(self.policy) + + # Fast path selection and scheduling (treat all as diffusion requests, + # independent of pooling_params) + while self.waiting and token_budget > 0 and capacity > 0: + request = self.waiting.peek_request() + # Uniformly treat as diffusion. A feature flag can be added later + # via config or request tag. + + # Allocate all input tokens for the request in one shot + # (allocate 1 placeholder if zero) + required_tokens = max(getattr(request, "num_prompt_tokens", 0), 1) + if required_tokens > token_budget: + # Insufficient budget to process all inputs at once; + # stop fast path attempt + break + num_new_tokens = required_tokens + new_blocks = self.kv_cache_manager.allocate_slots( + request, + num_new_tokens, + num_lookahead_tokens=self.num_lookahead_tokens, + ) + if new_blocks is None: + # Allocation failed (e.g., VRAM pressure); stop fast path and + # fall back to default scheduling + # Put the current request back to the head of the waiting queue + # Note: the original queue order is preserved + break + + # Officially schedule this request + request = self.waiting.pop_request() + self.running.append(request) + request.status = RequestStatus.RUNNING + if self.log_stats: + request.record_event(EngineCoreEventType.SCHEDULED, scheduled_timestamp) + + req_to_new_blocks[request.request_id] = new_blocks + num_scheduled_tokens[request.request_id] = num_new_tokens + token_budget -= num_new_tokens + capacity -= 1 + scheduled_new_reqs.append(request) + + # Return skipped waiting requests + if skipped_waiting_requests: + self.waiting.prepend_requests(skipped_waiting_requests) + + # If fast path scheduled none, fall back to the original scheduling + if not num_scheduled_tokens: + return super().schedule() + + # Compute common prefix blocks (aligned with v1) + num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups) + if self.running: + any_request = self.running[0] + num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks( + any_request, len(self.running) + ) + + # Assemble SchedulerOutput + new_reqs_data = [ + OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids()) + for req in scheduled_new_reqs + ] + cached_reqs_data = self._make_cached_request_data( + scheduled_running_reqs, + scheduled_resumed_reqs, + num_scheduled_tokens, + scheduled_spec_decode_tokens, + req_to_new_blocks, + ) + scheduled_requests = scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs + structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask( + scheduled_requests, scheduled_spec_decode_tokens + ) + + total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) + scheduler_output = SchedulerOutput( + scheduled_new_reqs=new_reqs_data, + scheduled_cached_reqs=cached_reqs_data, + num_scheduled_tokens=num_scheduled_tokens, + total_num_scheduled_tokens=total_num_scheduled_tokens, + scheduled_spec_decode_tokens=scheduled_spec_decode_tokens, + scheduled_encoder_inputs=scheduled_encoder_inputs, + num_common_prefix_blocks=num_common_prefix_blocks, + finished_req_ids=self.finished_req_ids, + free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), + structured_output_request_ids=structured_output_request_ids, + grammar_bitmask=grammar_bitmask, + ) + + # KVTransfer: package metadata + if self.connector is not None: + meta = self.connector.build_connector_meta(scheduler_output) + scheduler_output.kv_connector_metadata = meta + + # Update internal state (advance num_computed_tokens, free encoder inputs, + # etc.) + self._update_after_schedule(scheduler_output) + return scheduler_output + + """ + Scheduler for the diffusion model. + This scheduler is modified to stop the request immediately for the diffusion model. + This is because the diffusion model can generate the final image/audio in one step. + Note: This is just a minimal modification to the original scheduler, + and there should be some further efforts to optimize the scheduler. + The original scheduler is still used for the AR model. + """ + + def update_from_output( + self, + scheduler_output: SchedulerOutput, + model_runner_output: OmniModelRunnerOutput, + ) -> dict[int, EngineCoreOutputs]: + """Update the scheduler state based on the model runner output. + + This method is modified to stop the request immediately for the diffusion model. + """ + sampled_token_ids = model_runner_output.sampled_token_ids + logprobs = model_runner_output.logprobs + prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict + num_scheduled_tokens = scheduler_output.num_scheduled_tokens + pooler_outputs = model_runner_output.pooler_output + num_nans_in_logits = model_runner_output.num_nans_in_logits + kv_connector_output = model_runner_output.kv_connector_output + + outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) + spec_decoding_stats: Optional[SpecDecodingStats] = None + kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None + + # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, + # the below loop can be a performance bottleneck. We should do our best + # to avoid expensive operations inside the loop. + stopped_running_reqs: set[Request] = set() + stopped_preempted_reqs: set[Request] = set() + for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): + assert num_tokens_scheduled > 0 + request = self.requests.get(req_id) + if request is None: + # The request is already finished. This can happen if the + # request is aborted while the model is executing it (e.g., + # in pipeline parallelism). + continue + + req_index = model_runner_output.req_id_to_index[req_id] + generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else [] + + scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id) + if scheduled_spec_token_ids: + num_draft_tokens = len(scheduled_spec_token_ids) + num_accepted = len(generated_token_ids) - 1 + num_rejected = num_draft_tokens - num_accepted + # num_computed_tokens represents the number of tokens + # processed in the current step, considering scheduled + # tokens and rejections. If some tokens are rejected, + # num_computed_tokens is decreased by the number of rejected + # tokens. + request.num_computed_tokens -= num_rejected + spec_decoding_stats = self.make_spec_decoding_stats( + spec_decoding_stats, + num_draft_tokens=num_draft_tokens, + num_accepted_tokens=num_accepted, + ) + + new_logprobs = None + new_token_ids = generated_token_ids + kv_transfer_params = None + status_before_stop = request.status + pooler_output = None + if pooler_outputs: + pooler_output = pooler_outputs[req_index] + + # Diffusion request: completes in one step; mark finished and free resources + request.status = RequestStatus.FINISHED_STOPPED + # Optional: set a stop_reason for front-end clarity + # (does not affect protocol) + request.stop_reason = request.stop_reason # or "generation_done" + kv_transfer_params = self._free_request(request) + if status_before_stop == RequestStatus.RUNNING: + stopped_running_reqs.add(request) + else: + stopped_preempted_reqs.add(request) + + # Extract sample logprobs if needed. + if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs: + # NOTE: once we support N tokens per step (spec decode), + # the outer lists can be of length > 1. + new_logprobs = logprobs.slice(req_index, req_index + 1) + + if new_token_ids and self.structured_output_manager.should_advance(request): + # NOTE: structured_output_request should not be None if + # use_structured_output, we have check above, so safe to ignore + # type warning + request.structured_output_request.grammar.accept_tokens( # type: ignore[union-attr] # noqa: E501 + req_id, new_token_ids + ) + + # spec_token_ids comes from the model runner output + if num_nans_in_logits is not None and req_id in num_nans_in_logits: + request.num_nans_in_logits = num_nans_in_logits[req_id] + + # Get prompt logprobs for this request. + prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) + # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput + pooling_output_dict = None + if pooler_output is not None: + # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type + # which expects Optional[dict[str, torch.Tensor]] + pooling_output_dict = {"model_outputs": pooler_output} + if new_token_ids or pooler_output is not None or kv_transfer_params: + # Add EngineCoreOutput for this Request. + outputs[request.client_index].append( + EngineCoreOutput( + request_id=req_id, + new_token_ids=new_token_ids, + finish_reason=request.get_finished_reason(), + new_logprobs=new_logprobs, + new_prompt_logprobs_tensors=prompt_logprobs_tensors, + pooling_output=pooling_output_dict, + stop_reason=request.stop_reason, + events=request.take_events(), + kv_transfer_params=kv_transfer_params, + trace_headers=request.trace_headers, + num_cached_tokens=request.num_cached_tokens, + ) + ) + else: + # Invariant: EngineCore returns no partial prefill outputs. + assert not prompt_logprobs_tensors + + # Remove the stopped requests from the running and waiting queues. + if stopped_running_reqs: + if stopped_running_reqs: + self.running = remove_all(self.running, stopped_running_reqs) + if stopped_preempted_reqs: + # This is a rare case and unlikely to impact performance. + self.waiting.remove_requests(stopped_preempted_reqs) + + # KV Connector: update state for finished KV Transfers. + if model_runner_output.kv_connector_output: + self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output) + + # Create EngineCoreOutputs for all clients that have requests with + # outputs in this step. + engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()} + + finished_req_ids = self.finished_req_ids_dict + if finished_req_ids: + # Include ids of requests that finished since last outputs + # were sent. + for client_index, finished_set in finished_req_ids.items(): + # Set finished request set in EngineCoreOutputs for this client. + if (eco := engine_core_outputs.get(client_index)) is not None: + eco.finished_requests = finished_set + else: + engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set) + finished_req_ids.clear() + + if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None: + # Return stats to only one of the front-ends. + if (eco := next(iter(engine_core_outputs.values()), None)) is None: + # We must return the stats even if there are no request + # outputs this step. + engine_core_outputs[0] = eco = EngineCoreOutputs() + eco.scheduler_stats = stats + + return engine_core_outputs diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 7cc4066c9..8ce5f043a 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -2,10 +2,21 @@ Engine components for vLLM-Omni. """ -from typing import Any, Optional +import time +from collections.abc import Mapping +from typing import Any, Optional, Union import msgspec -from vllm.v1.engine import EngineCoreRequest +import torch +from vllm.v1.engine import ( + EngineCoreEvent, + EngineCoreRequest, + FinishReason, + LogprobsLists, + LogprobsTensors, + SchedulerStats, + UtilityOutput, +) class PromptEmbedsPayload(msgspec.Struct): @@ -66,3 +77,62 @@ class OmniEngineCoreRequest(EngineCoreRequest): prompt_embeds: Optional[PromptEmbedsPayload] = None # Optional additional information dictionary (serialized) additional_information: Optional[AdditionalInformationPayload] = None + + +class OmniEngineCoreOutput( + msgspec.Struct, + array_like=True, # type: ignore[call-arg] + omit_defaults=True, # type: ignore[call-arg] + gc=False, +): # type: ignore[call-arg] + request_id: str + new_token_ids: list[int] + + new_logprobs: Optional[LogprobsLists] = None + new_prompt_logprobs_tensors: Optional[LogprobsTensors] = None + + pooling_output: Optional[dict[str, torch.Tensor]] = None + + finish_reason: Optional[FinishReason] = None + stop_reason: Union[int, str, None] = None + events: Optional[list[EngineCoreEvent]] = None + kv_transfer_params: Optional[dict[str, Any]] = None + + trace_headers: Optional[Mapping[str, str]] = None + # The number of tokens with prefix cache hits. + num_cached_tokens: int = 0 + + @property + def finished(self) -> bool: + return self.finish_reason is not None + + +class OmniEngineCoreOutputs( + msgspec.Struct, + array_like=True, # type: ignore[call-arg] + omit_defaults=True, # type: ignore[call-arg] + gc=False, +): # type: ignore[call-arg] + # NOTE(Nick): We could consider ways to make this more compact, + # e.g. columnwise layout + + engine_index: int = 0 + + # [num_reqs] + outputs: list[OmniEngineCoreOutput] = [] + scheduler_stats: Optional[SchedulerStats] = None + timestamp: float = 0.0 + + utility_output: Optional[UtilityOutput] = None + finished_requests: Optional[set[str]] = None + + # In DP case, used to signal that the current wave of requests + # has finished and the engines are paused. + wave_complete: Optional[int] = None + # In DP case, used to signal that a request was received for an + # "old" wave, so the next wave needs to be started in other engines. + start_wave: Optional[int] = None + + def __post_init__(self): + if self.timestamp == 0.0: + self.timestamp = time.monotonic() diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 6e7b02be4..fc6dd5d7b 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -1,19 +1,21 @@ from dataclasses import dataclass from typing import Optional +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import Qwen3OmniMoeTextConfig from vllm.engine.arg_utils import EngineArgs +from vllm.logger import init_logger from vllm.v1.engine.async_llm import AsyncEngineArgs from vllm_omni.config import OmniModelConfig +logger = init_logger(__name__) + @dataclass class OmniEngineArgs(EngineArgs): """Engine arguments for omni models, extending base EngineArgs. - Adds omni-specific configuration fields for multi-stage pipeline processing and output type specification. - Args: stage_id: Identifier for the stage in a multi-stage pipeline (default: 0) model_stage: Stage type identifier, e.g., "thinker" or "talker" @@ -29,13 +31,20 @@ class OmniEngineArgs(EngineArgs): model_stage: str = "thinker" model_arch: str = "Qwen2_5OmniForConditionalGeneration" engine_output_type: Optional[str] = None + hf_config_name: Optional[str] = None + + def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig: + # transformers' get_text_config method is used to get the text config from thinker_config. + # to handle the case that each model stage has their own text config, + # we need to draw the text config from the corresponding model stage. + return getattr(config_dict["hf_config"], config_dict["hf_config_name"]).get_text_config() def create_model_config(self) -> OmniModelConfig: """Create an OmniModelConfig from these engine arguments. - Returns: OmniModelConfig instance with all configuration fields set """ + # First, get the base ModelConfig from the parent class base_config = super().create_model_config() @@ -48,7 +57,9 @@ class OmniEngineArgs(EngineArgs): config_dict["model_stage"] = self.model_stage config_dict["model_arch"] = self.model_arch config_dict["engine_output_type"] = self.engine_output_type - + config_dict["hf_config_name"] = self.hf_config_name + if self.hf_config_name is not None: + config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict) # Create and return the OmniModelConfig instance omni_config = OmniModelConfig(**config_dict) omni_config.hf_config.architectures = omni_config.architectures @@ -59,10 +70,8 @@ class OmniEngineArgs(EngineArgs): @dataclass class AsyncOmniEngineArgs(AsyncEngineArgs): """Async engine arguments for omni models, extending base AsyncEngineArgs. - Adds omni-specific configuration fields for multi-stage pipeline processing and output type specification in async contexts. - Args: stage_id: Identifier for the stage in a multi-stage pipeline (default: 0) model_stage: Stage type identifier, e.g., "thinker" or "talker" @@ -78,13 +87,15 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): model_stage: str = "thinker" model_arch: str = "Qwen2_5OmniForConditionalGeneration" engine_output_type: Optional[str] = None + hf_config_name: Optional[str] = None + + def draw_hf_text_config(self, config_dict: dict) -> Qwen3OmniMoeTextConfig: + # transformers' get_text_config method is used to get the text config from thinker_config. + # to handle the case that each model stage has their own text config, + # we need to draw the text config from the corresponding model stage. + return getattr(config_dict["hf_config"], config_dict["hf_config_name"]).get_text_config() def create_model_config(self) -> OmniModelConfig: - """Create an OmniModelConfig from these engine arguments. - - Returns: - OmniModelConfig instance with all configuration fields set - """ # First, get the base ModelConfig from the parent class base_config = super().create_model_config() @@ -98,6 +109,9 @@ class AsyncOmniEngineArgs(AsyncEngineArgs): config_dict["model_arch"] = self.model_arch config_dict["engine_output_type"] = self.engine_output_type + config_dict["hf_config_name"] = self.hf_config_name + if self.hf_config_name is not None: + config_dict["hf_text_config"] = self.draw_hf_text_config(config_dict) # Create and return the OmniModelConfig instance omni_config = OmniModelConfig(**config_dict) omni_config.hf_config.architectures = omni_config.architectures diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 97bab2633..63e5ecdab 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -1,3 +1,4 @@ +from ast import Dict from typing import Any, Callable, Optional, Union import torch @@ -33,7 +34,7 @@ class OmniRequestState(RequestState): ): super().__init__(*args, **kwargs) self.mm_type: Optional[str] = None - self.mm_accumulated: Optional[torch.Tensor] = None + self.mm_accumulated: Optional[Dict[str, Any]] = None @classmethod def from_new_request( @@ -92,33 +93,71 @@ class OmniRequestState(RequestState): log_stats=log_stats, ) - def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], mm_type: Optional[str]) -> None: - """Add a multimodal tensor to the accumulated outputs. - - Accumulates multimodal tensors (e.g., image latents, audio waveforms) - that are produced incrementally during generation. Tensors are - concatenated along the first dimension. - - Args: - tensor: Optional tensor to add. If None, this is a no-op. - mm_type: Optional modality type identifier (e.g., "image", "audio"). - Used to categorize the output type. - """ - if tensor is None: + def add_multimodal_tensor(self, payload: Optional[Any], mm_type: Optional[str]) -> None: + if payload is None: return try: if mm_type: self.mm_type = (mm_type or "").lower() - t = tensor.detach() - try: - t = t.to("cpu") - except Exception: - # Best-effort CPU move; keep original device if conversion fails - logger.debug("Failed to move multimodal tensor to CPU", exc_info=True) - if self.mm_accumulated is None: - self.mm_accumulated = t + + # Normalize incoming payload to dict on CPU + def _to_cpu(x): + if isinstance(x, torch.Tensor): + try: + return x.detach().to("cpu").contiguous() + except Exception: + return x + return x + + if isinstance(payload, dict): + incoming: Dict[str, Any] = {} + # Optional remap: if producer used "model_outputs" or "hidden", rename to mm_type + # to keep a consistent key namespace per engine_core_output_type. + remapped_tensor = dict(payload) + target_key = self.mm_type or "hidden" + if "model_outputs" in remapped_tensor: + remapped_tensor[target_key] = remapped_tensor.pop("model_outputs") + elif "hidden" in remapped_tensor and target_key != "hidden": + remapped_tensor[target_key] = remapped_tensor.pop("hidden") + for k, v in remapped_tensor.items(): + if isinstance(v, dict): + incoming[k] = {str(sk): _to_cpu(sv) for sk, sv in v.items()} + else: + incoming[k] = _to_cpu(v) else: - self.mm_accumulated = torch.cat([self.mm_accumulated, t], dim=0) + key = self.mm_type or "hidden" + incoming = {key: _to_cpu(payload)} + + if self.mm_accumulated is None: + self.mm_accumulated = incoming + else: + # Merge keys; concatenate tensors along token dim when possible + for k, v in incoming.items(): + if k not in self.mm_accumulated: + self.mm_accumulated[k] = v + else: + existing = self.mm_accumulated[k] + if isinstance(v, torch.Tensor) and isinstance(existing, torch.Tensor): + try: + self.mm_accumulated[k] = torch.cat([existing, v], dim=0) # type: ignore[index] + except Exception: + self.mm_accumulated[k] = v + elif isinstance(v, dict) and isinstance(existing, dict): + # Merge nested dicts by concatenating tensors along token dim when possible + for sk, sv in v.items(): + if sk not in existing: + existing[sk] = sv + continue + ev = existing[sk] + if isinstance(sv, torch.Tensor) and isinstance(ev, torch.Tensor): + try: + existing[sk] = torch.cat([ev, sv], dim=0) + except Exception: + existing[sk] = sv + else: + existing[sk] = sv + else: + self.mm_accumulated[k] = v except Exception: # Log and continue without crashing the output pipeline logger.exception("Error accumulating multimodal tensor") @@ -178,18 +217,15 @@ class OmniRequestState(RequestState): base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason) try: if self.mm_accumulated is not None: - tensor = self.mm_accumulated - try: - tensor = tensor.detach().to("cpu") - except Exception: - logger.debug( - "Failed to move accumulated multimodal tensor to CPU", - exc_info=True, - ) - # Attach on the completion output for downstream consumers. + # Attach accumulated multimodal dict on the completion output if not hasattr(base_output, "multimodal_output"): setattr(base_output, "multimodal_output", {}) - setattr(base_output, "multimodal_output", {self.mm_type: tensor}) + mm_out = getattr(base_output, "multimodal_output") + if isinstance(mm_out, dict): + for k, v in self.mm_accumulated.items(): + mm_out[k] = v + else: + setattr(base_output, "multimodal_output", self.mm_accumulated) except Exception: logger.exception("Error in _new_completion_output") return base_output @@ -369,10 +405,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): # Attach accumulated multimodal payload if any try: if isinstance(req_state, OmniRequestState) and req_state.mm_accumulated is not None: - mm_key = req_state.mm_type or "latents" if not hasattr(ro, "multimodal_output"): setattr(ro, "multimodal_output", {}) - ro.multimodal_output[mm_key] = req_state.mm_accumulated + ro.multimodal_output = req_state.mm_accumulated except Exception: logger.exception("Error attaching multimodal payload in process_outputs") if req_state.queue is not None: diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py index e1b74ea6f..5046196d0 100644 --- a/vllm_omni/entrypoints/async_omni_llm.py +++ b/vllm_omni/entrypoints/async_omni_llm.py @@ -2,7 +2,6 @@ import asyncio import multiprocessing as mp import os import socket -import sys import time from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence from concurrent.futures import ThreadPoolExecutor, as_completed @@ -94,8 +93,8 @@ class AsyncOmniLLM(EngineClient): init_sleep_seconds: int = 20, shm_threshold_bytes: int = 65536, batch_timeout: int = 10, - init_timeout: int = 300, - **kwargs: Any, + init_timeout: int = 60000, + **kwargs, ): self.batch_timeout = batch_timeout self._enable_stats: bool = bool(log_stats) @@ -194,6 +193,8 @@ class AsyncOmniLLM(EngineClient): logger.warning("[Orchestrator] Failed to stop stage worker: %s", e) def __del__(self) -> None: # best-effort + print("[AsyncOmniLLM] __del__ close()", flush=True) + raise Exception("test") try: self.close() except Exception as e: @@ -460,9 +461,8 @@ class AsyncOmniLLM(EngineClient): logger.exception("[Orchestrator] Failed to build/log summary: %s", e) def _wait_for_stages_ready(self, timeout: int = 120) -> None: - deadline = time.time() + max(0, int(timeout)) num_stages = len(self.stage_list) - while len(self._stages_ready) < num_stages and time.time() < deadline: + while len(self._stages_ready) < num_stages: progressed = False for stage_id, stage in enumerate(self.stage_list): if stage_id in self._stages_ready: @@ -528,20 +528,6 @@ class AsyncOmniLLM(EngineClient): occurred while logging suggestions", ) - # Attempt graceful shutdown of all stages before exiting - try: - self.close() - except Exception: - pass - - # Terminate the current process with non-zero exit code - try: - sys.exit(1) - except SystemExit: - raise - except Exception: - os._exit(1) - @property def is_running(self) -> bool: # Is None before the loop is started. diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 0a04ab406..4b9b1b4f6 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,6 +1,5 @@ import multiprocessing as mp import os -import sys import time from collections.abc import Sequence from concurrent.futures import ThreadPoolExecutor, as_completed @@ -11,7 +10,6 @@ from pydantic import ValidationError # External library imports (vLLM) from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field -from vllm.engine.arg_utils import HfOverrides from vllm.entrypoints.llm import LLM from vllm.inputs import PromptType from vllm.logger import init_logger @@ -96,12 +94,12 @@ class OmniLLM: # Optional file handler for orchestrator self._log_file = log_file - if self._log_file: - remove_old_logs(self._log_file, len(self.stage_configs)) - configure_orchestrator_logger(logger, self._log_file) self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file) self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout) + if self._log_file: + remove_old_logs(self._log_file, len(self.stage_list)) + configure_orchestrator_logger(logger, self._log_file) def _initialize_stages( self, @@ -378,7 +376,16 @@ class OmniLLM: next_stage_id = stage_id + 1 if next_stage_id < num_stages: next_stage: OmniStage = self.stage_list[next_stage_id] - next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]]) + try: + next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]]) + except Exception as e: + logger.exception( + "[Orchestrator] Process engine inputs error for req %s at stage %s: %s", + req_id, + next_stage_id, + e, + ) + continue sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index] try: # Measure transfer size and time (encode + enqueue) @@ -467,7 +474,7 @@ class OmniLLM: progressed = True if result.get("type") == "stage_ready": self._stages_ready.add(stage_id) - logger.debug("[Orchestrator] Stage-%s reported ready", stage_id) + logger.info("[Orchestrator] Stage-%s reported ready", stage_id) else: # No user data should arrive before seeding; ignore other messages pass @@ -500,20 +507,8 @@ class OmniLLM: logger.error( "[Orchestrator] Stage initialization failed and an error occurred while logging suggestions", ) - - # Attempt graceful shutdown of all stages before exiting - try: - self.close() - except Exception: - pass - - # Terminate the current process with non-zero exit code - try: - sys.exit(1) - except SystemExit: - raise - except Exception: - os._exit(1) + elif len(self._stages_ready) == num_stages: + logger.info("[Orchestrator] All stages initialized successfully") class OmniStageLLM(LLM): @@ -539,9 +534,8 @@ class OmniStageLLM(LLM): self, model: str, compilation_config: Optional[Union[int, dict[str, Any], CompilationConfig]] = None, - hf_overrides: Optional[HfOverrides] = None, structured_outputs_config: Optional[Union[dict[str, Any], StructuredOutputsConfig]] = None, - **kwargs: Any, + **kwargs, ): """LLM constructor.""" if "disable_log_stats" not in kwargs: @@ -570,9 +564,6 @@ class OmniStageLLM(LLM): # to provide better context to the user. raise ValueError(f"Invalid 'kv_transfer_config' provided: {e}") from e - if hf_overrides is None: - hf_overrides = {} - if compilation_config is not None: if isinstance(compilation_config, int): compilation_config_instance = CompilationConfig(level=compilation_config) @@ -597,7 +588,6 @@ class OmniStageLLM(LLM): engine_args = OmniEngineArgs( model=model, - hf_overrides=hf_overrides, compilation_config=compilation_config_instance, structured_outputs_config=structured_outputs_instance, **kwargs, diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 624c943fe..d44c0043b 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -14,6 +14,7 @@ import asyncio import importlib import logging import multiprocessing as mp +import os from typing import Any, Optional, Union from vllm.inputs import TextPrompt @@ -237,6 +238,14 @@ class OmniStage: self._proc.terminate() except Exception as e: self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e) + # Cleanup temporary stage log if we created one (only when no log_file provided) + try: + if not self._log_file: + tmp_path = f"/tmp/omni_stage{self.stage_id}.log" + if os.path.exists(tmp_path): + os.remove(tmp_path) + except Exception: + pass def submit(self, payload: dict[str, Any]) -> None: """Submit a task to the stage worker. @@ -286,8 +295,6 @@ class OmniStage: raise ValueError("engine_input_source is empty") source_stage_id = self.engine_input_source[0] source_outputs = stage_list[source_stage_id].engine_outputs - if not isinstance(prompt, list): - prompt = [prompt] multi_modal_data = { source_output.request_id: p.get("multi_modal_data", None) for source_output, p in zip(source_outputs, prompt) @@ -322,6 +329,7 @@ def _stage_worker( ) -> None: """Stage worker entry: device setup, LLM init, batching, SHM IPC.""" import logging as _logging + import os as _os import time as _time from vllm_omni.entrypoints.log_utils import ( @@ -331,7 +339,6 @@ def _stage_worker( log_stage_running_avg, ) from vllm_omni.entrypoints.omni_llm import OmniStageLLM - # no inline JSONL/serialization imports; logging handled by utilities stage_id = stage_payload["stage_id"] @@ -339,12 +346,31 @@ def _stage_worker( runtime_cfg = stage_payload.get("runtime", {}) shm_threshold_bytes = int(stage_payload.get("shm_threshold_bytes", 65536)) - # Per-stage file logger (optional) + # Per-stage logger: clear inherited handlers to avoid broken parent streams try: - if log_file: - stage_log = _logging.getLogger(__name__) - stage_log.setLevel(_logging.DEBUG) - fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log") + stage_log = _logging.getLogger(__name__) + stage_log.setLevel(_logging.DEBUG) + for _h in list(stage_log.handlers): + try: + stage_log.removeHandler(_h) + except Exception: + pass + stage_log.propagate = False + + class _StageFilter(_logging.Filter): + def filter(self, record: _logging.LogRecord) -> bool: + setattr(record, "stage", stage_id) + return True + + # Prefer file logging even when log_file is not provided, to avoid invalid stdio in child procs + try: + if log_file: + _path = f"{log_file}.stage{stage_id}.log" + else: + _path = f"/tmp/omni_stage{stage_id}.log" + # Ensure dir exists + _os.makedirs(_os.path.dirname(_path), exist_ok=True) + fh = _logging.FileHandler(_path) fh.setLevel(_logging.DEBUG) fh.setFormatter( _logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s") @@ -357,6 +383,34 @@ def _stage_worker( fh.addFilter(_StageFilter()) stage_log.addHandler(fh) + try: + # Also route vLLM internal logs to the same handler for this stage + _vllm_logger = _logging.getLogger("vllm") + _vllm_logger.setLevel(_logging.DEBUG) + for _vh in list(_vllm_logger.handlers): + try: + _vllm_logger.removeHandler(_vh) + except Exception: + pass + _vllm_logger.propagate = False + _vllm_logger.addHandler(fh) + except Exception: + pass + except Exception: + # Final fallback: attach NullHandler to avoid logging errors + stage_log.addHandler(_logging.NullHandler()) + except Exception: + pass + + try: + _logging.raiseExceptions = False + _root_logger = _logging.getLogger() + for _h in list(_root_logger.handlers): + try: + _root_logger.removeHandler(_h) + except Exception: + pass + _root_logger.addHandler(_logging.NullHandler()) except Exception: pass @@ -377,9 +431,7 @@ def _stage_worker( # Init LLM _logging.getLogger(__name__).debug( - "[Stage-%s] Initializing engine with args keys=%s", - stage_id, - list(engine_args.keys()), + "[Stage-%s] Initializing engine with args keys=%s", stage_id, list(engine_args.keys()) ) stage_engine = OmniStageLLM(model=model, **engine_args) _logging.getLogger(__name__).debug("[Stage-%s] Engine initialized", stage_id) @@ -394,21 +446,35 @@ def _stage_worker( task = in_q.get() _recv_dequeue_ts = _time.time() if task is None: - _logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id) + _logging.getLogger(__name__).error("[Stage-%s] Received shutdown signal", stage_id) break max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1) + print(f"[Stage-{stage_id}] Max batch size: {max_batch_size}") batch_tasks: list[dict[str, Any]] = [task] if max_batch_size > 1: while len(batch_tasks) < max_batch_size: + start_time = _time.time() if not in_q.empty(): - extra = in_q.get(timeout=batch_timeout) + extra = in_q.get_nowait() if extra is None: in_q.put(None) break batch_tasks.append(extra) + end_time = _time.time() + duration = end_time - start_time + if duration > batch_timeout: + break + else: + continue else: - break + end_time = _time.time() + duration = end_time - start_time + _time.sleep(0.05) + if duration > batch_timeout: + break + else: + continue batch_request_ids: list[Any] = [] batch_engine_inputs: list[Any] = [] @@ -434,7 +500,7 @@ def _stage_worker( elif isinstance(ein, dict): batch_engine_inputs.append(ein) else: - _logging.getLogger(__name__).exception("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein)) + _logging.getLogger(__name__).error("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein)) sampling_params = batch_tasks[0]["sampling_params"] _logging.getLogger(__name__).debug( "[Stage-%s] Received batch size=%d, request_ids=%s", @@ -456,6 +522,14 @@ def _stage_worker( gen_outputs.append(ro) _gen_t1 = _time.time() _gen_ms = (_gen_t1 - _gen_t0) * 1000.0 + try: + print( + f"[Stage-{stage_id}] Generate done: batch={len(batch_tasks)}, " + f"req_ids={batch_request_ids}, gen_ms={_gen_ms:.1f}", + flush=True, + ) + except Exception: + pass # Group outputs per request id with fallback req_to_outputs: dict[Any, list[Any]] = {rid: [] for rid in batch_request_ids} @@ -783,6 +857,14 @@ async def _stage_worker_async( "metrics": _metrics, } ) + try: + print( + f"[Stage-{stage_id}] Enqueued req={rid}, use_shm={use_shm}, " + f"tokens_out={_metrics['num_tokens_out']}", + flush=True, + ) + except Exception: + pass except Exception as e: _logging.getLogger(__name__).exception( "[Stage-%s] Failed to enqueue result for request %s: %s", @@ -815,3 +897,6 @@ async def _stage_worker_async( "error": str(e), } ) + print("--------------------------------", flush=True) + print(f"[Stage-{stage_id}] Stage worker exiting", flush=True) + print("--------------------------------", flush=True) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 14c266c85..0cd522e14 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -730,7 +730,7 @@ class OmniOpenAIServingChat(OpenAIServingChat): def _create_audio_choice(self, omni_outputs: OmniRequestOutput, role: str): choices: list[ChatCompletionResponseChoice] = [] final_res = omni_outputs.request_output - audio_tensor = final_res.multimodal_output["audio"].detach().cpu().numpy() + audio_tensor = final_res.multimodal_output["audio"].float().detach().cpu().numpy() # Convert numpy array to WAV bytes and encode as base64 if soundfile is None: diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index 73179259a..8ddc7d504 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -53,11 +53,7 @@ def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None: if 0 <= logical_idx < len(mapping): selected_physical = mapping[logical_idx] except Exception as e: - logger.debug( - "[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", - stage_id, - e, - ) + logger.debug("[Stage-%s] Failed to map logical index via CUDA_VISIBLE_DEVICES: %s", stage_id, e) selected_physical = None if selected_physical is None: selected_physical = int(logical_idx) @@ -69,18 +65,12 @@ def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None: selected_physical, ) elif devices in (None, "cpu"): - logger.debug( - "[Stage-%s] Using default device visibility (devices=%s)", - stage_id, - devices, - ) + logger.debug("[Stage-%s] Using default device visibility (devices=%s)", stage_id, devices) else: selected_physical = int(str(devices)) os.environ["CUDA_VISIBLE_DEVICES"] = str(selected_physical) logger.debug( - "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", - stage_id, - selected_physical, + "[Stage-%s] Set CUDA_VISIBLE_DEVICES to single device %s (fallback)", stage_id, selected_physical ) try: diff --git a/vllm_omni/model_executor/layers/mrope.py b/vllm_omni/model_executor/layers/mrope.py index 789b3b3cf..463653470 100644 --- a/vllm_omni/model_executor/layers/mrope.py +++ b/vllm_omni/model_executor/layers/mrope.py @@ -4,8 +4,11 @@ from typing import Optional, Union import numpy as np import torch from transformers import PretrainedConfig +from vllm.logger import init_logger from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding +logger = init_logger(__name__) + def _apply_rotary_emb( x: torch.Tensor, @@ -442,9 +445,16 @@ class MRotaryEmbedding(RotaryEmbedding): # _vl_get_input_positions_tensor. thinker_config = hf_config.thinker_config - audio_token_id = thinker_config.audio_token_index - image_token_id = thinker_config.image_token_index - video_token_id = thinker_config.video_token_index + try: + audio_token_id = thinker_config.audio_token_index + image_token_id = thinker_config.image_token_index + video_token_id = thinker_config.video_token_index + except Exception: + logger.info("Multimodal token idx changed!") + audio_token_id = thinker_config.audio_token_id + image_token_id = thinker_config.image_token_id + video_token_id = thinker_config.video_token_id + audio_start_token_id = thinker_config.audio_start_token_id audio_end_token_id = thinker_config.audio_end_token_id vision_start_token_id = thinker_config.vision_start_token_id diff --git a/vllm_omni/model_executor/models/__init__.py b/vllm_omni/model_executor/models/__init__.py index 46b0e61e0..0b2629b4a 100644 --- a/vllm_omni/model_executor/models/__init__.py +++ b/vllm_omni/model_executor/models/__init__.py @@ -1 +1,4 @@ +from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration from .registry import OmniModelRegistry # noqa: F401 + +__all__ = ["Qwen3OmniMoeForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/output_templates.py b/vllm_omni/model_executor/models/output_templates.py new file mode 100644 index 000000000..6fcd53eb3 --- /dev/null +++ b/vllm_omni/model_executor/models/output_templates.py @@ -0,0 +1,13 @@ +from typing import NamedTuple, Optional + +import torch +from vllm.sequence import IntermediateTensors + + +class OmniOutput(NamedTuple): + """Output from the merged Omni model containing both text and audio.""" + + text_hidden_states: torch.Tensor + multimodal_outputs: Optional[dict] = None + intermediate_tensors: Optional[IntermediateTensors] = None + next_token_id: Optional[torch.Tensor] = None diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index b650712b5..34961441f 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -2,7 +2,7 @@ import glob import os from collections.abc import Iterable from functools import cached_property -from typing import NamedTuple, Optional, Union +from typing import Optional, Union import numpy as np import torch @@ -26,6 +26,7 @@ from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific +from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( Qwen2_5OmniConditionalGenerationMixin, Qwen2_5OmniThinkerDummyInputsBuilder, @@ -39,14 +40,6 @@ TALKER_CODEC_EOS_TOKEN_ID = 8294 TALKER_CODEC_BOS_TOKEN_ID = 8293 -class OmniOutput(NamedTuple): - """Output from the merged Omni model containing both text and audio.""" - - text_hidden_states: torch.Tensor - multimodal_outputs: dict = {} - intermediate_tensors: Optional[IntermediateTensors] = None - - logger = init_logger(__name__) @@ -439,7 +432,7 @@ class Qwen2_5OmniForConditionalGeneration( code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code audio_tensor = self.generate_audio(code, voice_type) - return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audio_tensor}) + return OmniOutput(text_hidden_states=None, multimodal_outputs={"model_outputs": audio_tensor}) return OmniOutput( text_hidden_states=torch.cat( diff --git a/vllm_omni/model_executor/models/qwen3_omni/__init__.py b/vllm_omni/model_executor/models/qwen3_omni/__init__.py new file mode 100644 index 000000000..a469704f6 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/__init__.py @@ -0,0 +1,3 @@ +from .qwen3_omni import Qwen3OmniMoeForConditionalGeneration + +__all__ = ["Qwen3OmniMoeForConditionalGeneration"] diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py new file mode 100644 index 000000000..933236313 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py @@ -0,0 +1,173 @@ +from __future__ import annotations + +import torch +import torch.nn.functional as F +from torch import nn +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.models.qwen3_moe import ( + Qwen3MoeDecoderLayer, + Qwen3MoeMLP, + Qwen3MoeModel, # as _BaseQwen3MoeModel, +) +from vllm.model_executor.models.qwen3_moe import ( + Qwen3MoeForCausalLM as _BaseQwen3MoeForCausalLM, +) +from vllm.model_executor.models.utils import ( + PPMissingLayer, + maybe_prefix, +) + +logger = init_logger(__name__) + + +# Individual expert MoE block using Qwen3MoeMLP instead of FusedMoE +class Qwen3OmniMoeSparseMoeBlock(nn.Module): + """Sparse MoE block using individual Qwen3MoeMLP experts instead of FusedMoE.""" + + def __init__( + self, + vllm_config: VllmConfig, + prefix: str = "", + ): + super().__init__() + + config = vllm_config.model_config.hf_text_config + quant_config = vllm_config.quant_config + + self.num_experts = config.num_experts + self.top_k = config.num_experts_per_tok + self.norm_topk_prob = config.norm_topk_prob + self.hidden_size = config.hidden_size + + # Create individual expert MLPs + self.experts = nn.ModuleList( + [ + Qwen3MoeMLP( + hidden_size=config.hidden_size, + intermediate_size=config.moe_intermediate_size, + hidden_act=config.hidden_act, + quant_config=quant_config, + reduce_results=False, + prefix=f"{prefix}.experts.{i}", + ) + for i in range(self.num_experts) + ] + ) + + # Router for expert selection + from vllm.model_executor.layers.linear import ReplicatedLinear + + self.gate = ReplicatedLinear( + config.hidden_size, config.num_experts, bias=False, quant_config=quant_config, prefix=f"{prefix}.gate" + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + """Forward pass using individual experts.""" + # Handle 3D inputs (batch, seq_len, hidden_size) by reshaping to 2D + orig_shape = hidden_states.shape + if hidden_states.dim() == 3: + batch_size, seq_len, hidden_dim = hidden_states.shape + hidden_states = hidden_states.view(-1, hidden_dim) + elif hidden_states.dim() == 2: + num_tokens, hidden_dim = hidden_states.shape + elif hidden_states.dim() == 1: + hidden_states = hidden_states.unsqueeze(0) + num_tokens, hidden_dim = hidden_states.shape + else: + raise ValueError( + f"Qwen3OmniMoeSparseMoeBlock only supports 1D, 2D, or 3D inputs, got {hidden_states.dim()}D" + ) + + is_input_1d = len(orig_shape) == 1 + hidden_states = hidden_states.view(-1, hidden_dim) + + # Get router logits and select experts (matching transformers) + router_logits, _ = self.gate(hidden_states) + selected_experts, routing_weights = self._route_tokens(router_logits) + + # Forward through individual experts + final_hidden_states = self._forward_experts(hidden_states, selected_experts, routing_weights) + + # Reshape back to original shape + if is_input_1d: + return final_hidden_states.squeeze(0) + elif len(orig_shape) == 3: + # Reshape back to 3D (batch, seq_len, hidden_dim) + return final_hidden_states.view(orig_shape) + else: + return final_hidden_states + + def _route_tokens(self, router_logits: torch.Tensor): + """Route tokens to experts using top-k selection (matching transformers).""" + routing_weights = F.softmax(router_logits, dim=-1, dtype=torch.float) + routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1) + if self.norm_topk_prob: + routing_weights /= routing_weights.sum(dim=-1, keepdim=True) + routing_weights = routing_weights.to(router_logits.dtype) + return selected_experts, routing_weights + + def _forward_experts( + self, hidden_states: torch.Tensor, selected_experts: torch.Tensor, routing_weights: torch.Tensor + ): + """Forward through individual experts (matching transformers implementation).""" + final_hidden_states = torch.zeros_like(hidden_states) + expert_mask = torch.nn.functional.one_hot(selected_experts, num_classes=self.num_experts).permute(2, 1, 0) + + expert_hit = torch.greater(expert_mask.sum(dim=(-1, -2)), 0).nonzero() + for expert_idx in expert_hit: + idx, top_x = torch.where(expert_mask[expert_idx].squeeze(0)) + current_state = hidden_states[None, top_x].reshape(-1, hidden_states.shape[-1]) + current_hidden_states = self.experts[expert_idx](current_state) * routing_weights[top_x, idx, None] + final_hidden_states.index_add_(0, top_x, current_hidden_states.to(hidden_states.dtype)) + + return final_hidden_states + + +class Qwen3MoeForCausalLM(_BaseQwen3MoeForCausalLM): + """Thin wrapper to swap in the patched `Qwen3MoeModel`.""" + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + # Don't call super().__init__() to avoid duplicate layer registration. + nn.Module.__init__(self) + config = vllm_config.model_config.hf_text_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + self.model = Qwen3MoeModel(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + self.lm_head = ParallelLMHead( + config.vocab_size, config.hidden_size, quant_config=quant_config, prefix=maybe_prefix(prefix, "lm_head") + ) + if self.config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + # Set MoE hyperparameters for individual experts + self.expert_weights = [] + + self.moe_layers: list[FusedMoE] = [] + example_layer = None + for layer in self.model.layers: + if isinstance(layer, PPMissingLayer): + continue + + assert isinstance(layer, Qwen3MoeDecoderLayer) + if isinstance(layer.mlp, FusedMoE): + example_layer = layer.mlp + self.moe_layers.append(layer.mlp) + + if example_layer is None: + raise RuntimeError("No Qwen3OmniMoe layer found in the model.layers.") + + self.num_moe_layers = len(self.moe_layers) + self.num_expert_groups = 1 + self.num_shared_experts = 0 + self.num_logical_experts = example_layer.n_logical_experts + self.num_physical_experts = example_layer.n_physical_experts + self.num_local_physical_experts = example_layer.n_local_physical_experts + self.num_routed_experts = example_layer.n_routed_experts + self.num_redundant_experts = example_layer.n_redundant_experts diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py new file mode 100644 index 000000000..f38523600 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -0,0 +1,1132 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2025 The Qwen team. +"""Inference-only Qwen3-Omni-Moe unified model (thinker + talker + code2wav).""" + +from collections import defaultdict +from collections.abc import Iterable +from functools import cached_property +from typing import Optional, Union + +import torch +import torch.nn as nn +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeCode2WavConfig, + Qwen3OmniMoeConfig, + Qwen3OmniMoeTalkerConfig, + Qwen3OmniMoeThinkerConfig, +) +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP +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.v1.outputs import SamplerOutput +from vllm.v1.sample.metadata import SamplingMetadata +from vllm.v1.sample.sampler import Sampler + +from vllm_omni.model_executor.models.output_templates import OmniOutput +from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights + +from .qwen3_omni_moe_thinker import ( + Qwen3OmniMoeConditionalGenerationMixin, + Qwen3OmniMoeThinkerDummyInputsBuilder, + Qwen3OmniMoeThinkerMultiModalProcessor, + Qwen3OmniMoeThinkerProcessingInfo, +) + +# Special token IDs for Qwen3 Omni MoE +# Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json + +# Audio tokens (thinker vocabulary, for marking audio boundaries) +AUDIO_START_TOKEN_ID = 151669 # <|audio_start|> (audio_bos_token) +AUDIO_END_TOKEN_ID = 151670 # <|audio_end|> (audio_eos_token) +AUDIO_PAD_TOKEN_ID = 151675 # <|audio_pad|> + +# TTS text tokens (thinker vocabulary, for text-to-speech control) +TTS_PAD_TOKEN_ID = 151671 # +TTS_BOS_TOKEN_ID = 151672 # +TTS_EOS_TOKEN_ID = 151673 # (end of dialogue) +TTS_BOS_SINGLE_TOKEN_ID = 151674 # + +# Talker codec tokens (talker vocabulary, used for RVQ code generation) +TALKER_CODEC_PAD_TOKEN_ID = 4196 # Padding token +TALKER_CODEC_BOS_TOKEN_ID = 4197 # Beginning of speech +TALKER_CODEC_EOS_TOKEN_ID = 4198 # End of speech +TALKER_CODEC_NOTHING_ID = 4203 # No-think mode +TALKER_CODEC_THINK_BOS_ID = 4204 # Think mode start +TALKER_CODEC_THINK_EOS_ID = 4205 # Think mode end + +logger = init_logger(__name__) + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3OmniMoeThinkerMultiModalProcessor, + info=Qwen3OmniMoeThinkerProcessingInfo, + dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder, +) +class Qwen3OmniMoeForConditionalGeneration( + nn.Module, SupportsMultiModal, SupportsPP, Qwen3OmniMoeConditionalGenerationMixin +): + """ + Unified Qwen3 Omni MoE model combining thinker, talker, and code2wav. + + Architecture: + - Thinker: Multimodal understanding (text + audio + video) → text generation + - Talker: Text embeddings → RVQ codec codes + - Code2Wav: RVQ codes → audio waveform + + Usage: + Set `model_stage` in vllm_config to one of: "thinker", "talker", "code2wav" + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + self.have_multimodal_outputs = True + config: Qwen3OmniMoeConfig = vllm_config.model_config.hf_config + multimodal_config = vllm_config.model_config.multimodal_config + + # Keep vllm_config for later submodule init + self.vllm_config = vllm_config + self.config = config + + # Initialize thinker components + thinker_config: Qwen3OmniMoeThinkerConfig = config.thinker_config + self.thinker_config = thinker_config + self.multimodal_config = multimodal_config + + # Initialize talker components + talker_config: Qwen3OmniMoeTalkerConfig = config.talker_config + self.talker_config = talker_config + + # Initialize code2wav components + code2wav_config: Qwen3OmniMoeCode2WavConfig = config.code2wav_config + self.code2wav_config = code2wav_config + + # Determine model stage + self.model_stage = vllm_config.model_config.model_stage + + if self.model_stage == "thinker": + # Initialize thinker model (multimodal processing + text generation) + # Create a new vllm_config with thinker_config as the hf_config + thinker_vllm_config = vllm_config.with_hf_config( + thinker_config, architectures=["Qwen3OmniMoeThinkerForConditionalGeneration"] + ) + self.thinker = init_vllm_registered_model( + vllm_config=thinker_vllm_config, + prefix=maybe_prefix(prefix, "thinker"), + hf_config=thinker_config, + architectures=["Qwen3OmniMoeThinkerForConditionalGeneration"], + ) + self.model = self.thinker + self.talker = None + self.code2wav = None + elif self.model_stage == "talker": + self.thinker = None + # Initialize talker model (text embeddings → codec codes) + # Create a new vllm_config with talker_config as the hf_config + # This ensures the talker uses its own text_config (smaller vocab_size) + talker_vllm_config = vllm_config.with_hf_config( + talker_config, architectures=["Qwen3OmniMoeTalkerForConditionalGeneration"] + ) + self.talker = init_vllm_registered_model( + vllm_config=talker_vllm_config, + prefix=maybe_prefix(prefix, "talker"), + hf_config=talker_config, + architectures=["Qwen3OmniMoeTalkerForConditionalGeneration"], + ) + self.talker.init_multi_modal(thinker_config) + self.model = self.talker + self.code2wav = None + from transformers.generation.logits_process import ( + LogitsProcessorList, + RepetitionPenaltyLogitsProcessor, + SuppressTokensLogitsProcessor, + TemperatureLogitsWarper, + TopKLogitsWarper, + ) + + self.logits_processors = LogitsProcessorList( + [ + RepetitionPenaltyLogitsProcessor(penalty=1.05), + TemperatureLogitsWarper(temperature=0.9), + SuppressTokensLogitsProcessor(suppress_tokens=self._get_talker_suppressed_tokens()), + TopKLogitsWarper(top_k=50), + ] + ) + + elif self.model_stage == "code2wav": + self.thinker = None + self.talker = None + # Initialize code2wav (codec codes → audio waveform) + # Create a new vllm_config with code2wav_config as the hf_config + code2wav_vllm_config = vllm_config.with_hf_config(code2wav_config, architectures=["Qwen3OmniMoeCode2Wav"]) + self.code2wav = init_vllm_registered_model( + vllm_config=code2wav_vllm_config, + prefix=maybe_prefix(prefix, "code2wav"), + hf_config=code2wav_config, + architectures=["Qwen3OmniMoeCode2Wav"], + ) + self.model = self.code2wav + else: + raise ValueError( + f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'talker', 'code2wav'" + ) + + # Set up intermediate tensors + self.make_empty_intermediate_tensors = ( + self.thinker.make_empty_intermediate_tensors if self.model_stage == "thinker" else lambda: None + ) + + # ==================== Device utilities ==================== + + @staticmethod + def _module_device(module: nn.Module) -> torch.device: + """Get the device of a module.""" + try: + return next(module.parameters()).device + except StopIteration: + # No parameters; fall back to buffers or cpu + for _, buf in module.named_buffers(recurse=True): + return buf.device + return torch.device("cpu") + + def move_submodules_to_devices( + self, + *, + thinker_device: Optional[Union[str, torch.device]] = None, + talker_device: Optional[Union[str, torch.device]] = None, + code_predictor_device: Optional[Union[str, torch.device]] = None, + code2wav_device: Optional[Union[str, torch.device]] = None, + ) -> None: + """ + Optionally move thinker/talker/code2wav to different devices. + + Example: + model.move_submodules_to_devices( + thinker_device='cuda:0', + talker_device='cuda:1', + code2wav_device='cuda:2', + ) + """ + if thinker_device is not None and self.thinker is not None: + self.thinker.to(thinker_device) + if talker_device is not None and self.talker is not None: + self.talker.to(talker_device) + if code2wav_device is not None and self.code2wav is not None: + self.code2wav.to(code2wav_device) + + @cached_property + def sampler(self): + """Get sampler from active model.""" + if hasattr(self.model, "sampler"): + return self.model.sampler + return Sampler() + + def get_input_embeddings( + self, + input_ids: torch.Tensor, + multimodal_embeddings=None, + ) -> torch.Tensor: + """Get input embeddings for the active model stage.""" + if self.model_stage == "code2wav": + # Code2wav doesn't use text embeddings + return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size()) + return self.model.get_input_embeddings(input_ids, multimodal_embeddings) + + def get_multimodal_embeddings(self, **kwargs): + """Delegate to active model for multimodal processing.""" + return self.model.get_multimodal_embeddings(**kwargs) + + # ==================== Forward Pass ==================== + def _get_talker_suppressed_tokens(self): + return [ + i + for i in range( + self.config.talker_config.text_config.vocab_size - 1024, + self.config.talker_config.text_config.vocab_size, + ) + if i != self.config.talker_config.codec_eos_token_id + ] + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: Optional[IntermediateTensors] = None, + inputs_embeds: Optional[torch.Tensor] = None, + generate_audio: bool = True, + voice_type: str = "ethan", + codec: Optional[torch.Tensor] = None, + sampling_metadata: Optional[SamplingMetadata] = None, + logits_index: Optional[int] = None, + additional_information: Optional[dict[str, object]] = None, + **kwargs: object, + ) -> Union[torch.Tensor, IntermediateTensors, OmniOutput]: + """ + Unified forward pass for all model stages. + + Workflow: + 1) Thinker: multimodal understanding → text hidden states + 2) Talker -> Code Predictor: text embeddings → codec codes (layer 0 + code_predictor:residual layers) + 3) Code2wav: 8-layer RVQ codes → audio waveform + + Returns: + OmniOutput with text_hidden_states and optional audio + """ + + # ========== Stage 1: Thinker ========== + if self.model_stage == "thinker": + # Normalize to batched inputs if needed + added_batch_dim = False + if input_ids is not None and input_ids.ndim == 1: + input_ids = input_ids.unsqueeze(0) + added_batch_dim = True + if positions is not None and positions.ndim == 1: + positions = positions.unsqueeze(0) + added_batch_dim = True + if inputs_embeds is not None and inputs_embeds.ndim == 2: + inputs_embeds = inputs_embeds.unsqueeze(0) + added_batch_dim = True + + thinker_dev = self._module_device(self.thinker) + + # Handle None input_ids + if input_ids is None: + input_ids = torch.zeros( + inputs_embeds.shape[1], + dtype=torch.long, + device=thinker_dev, + ).unsqueeze(0) + added_batch_dim = True + + # Move to thinker device + if input_ids is not None and input_ids.device != thinker_dev: + input_ids = input_ids.to(thinker_dev) + if positions is not None and positions.device != thinker_dev: + positions = positions.to(thinker_dev) + if inputs_embeds is not None and inputs_embeds.device != thinker_dev: + inputs_embeds = inputs_embeds.to(thinker_dev) + + # Run thinker forward + # If talker expects a specific intermediate layer, capture it here + accept_layer = getattr(self.talker_config, "accept_hidden_layer", None) + capture_kwargs = {} + if accept_layer is not None: + capture_kwargs = { + "capture_layer_indices": [0, int(accept_layer)], + "return_hidden_states": True, + } + text_hidden_states, captured_layer_dict = self.thinker( + input_ids=input_ids, + positions=positions[0] if positions.ndim > 1 else positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **capture_kwargs, + **kwargs, + ) + # Compute thinker-side TTS token embeddings for BOS/EOS/PAD and expose via multimodal outputs. + # These will later be projected into talker text space by the talker stage. + multimodal_outputs = captured_layer_dict if captured_layer_dict is not None else {} + try: + tts_tokens = torch.tensor( + [[self.config.tts_bos_token_id, self.config.tts_eos_token_id, self.config.tts_pad_token_id]], + device=self._module_device(self.thinker), + dtype=torch.long, + ) + thinker_tts_embeds = self.thinker.get_input_embeddings(tts_tokens) # [1,3,thinker_hidden] + if ( + isinstance(thinker_tts_embeds, torch.Tensor) + and thinker_tts_embeds.ndim == 3 + and thinker_tts_embeds.shape[1] == 3 + ): + bos_eos_pad = thinker_tts_embeds.to(text_hidden_states.device).chunk(3, dim=1) # 3 * [1,1,H] + multimodal_outputs["tts_bos_embed"] = bos_eos_pad[0] + multimodal_outputs["tts_eos_embed"] = bos_eos_pad[1] + multimodal_outputs["tts_pad_embed"] = bos_eos_pad[2] + except Exception: + # Best-effort; absence will be handled by talker with fallbacks + pass + + # Return text-only output (with multimodal sidecar) + return OmniOutput( + text_hidden_states=(text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states), + multimodal_outputs=multimodal_outputs, + ) + + # ========== Stage 2.1: Talker ========== + elif self.model_stage == "talker": + # Mixed-mode support: In a single step, both Prefill*n and Decode*n are supported. + # Rules: + # - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS] + # - Decode segments consist of a single non-special token. + # - If additional_information is provided (can be a list split by request + # or a concatenated tensor plus a list of shapes), then for each request, + # reconstruct the thinker→talker input embeddings for the Prefill segments; + # - For Decode segments, if per-request auxiliary decode embeddings are + # provided (optional), add them; otherwise, keep the original embedding. + + if input_ids is None and additional_information is None: + input_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device) + additional_information = {} + self.code_predictor_hidden_pool = torch.zeros_like(inputs_embeds) + is_profile = True + else: + is_profile = False + + # Ensure we have base embeddings when only ids are provided + if inputs_embeds is None and input_ids is not None: + inputs_embeds = self.talker.get_input_embeddings(input_ids) + + # ------- Request-scoped additional information (no cross-request concat) ------- + request_ids: Optional[list[str]] = kwargs.get("request_ids") # ordered + request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get("request_token_spans") + addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id") + runtime_addi = kwargs.get("runtime_additional_information") + + # Normalize runtime_addi into a mapping by request_id for convenience + runtime_addi_by_req: dict[str, dict] = {} + if ( + isinstance(request_ids, list) + and isinstance(runtime_addi, list) + and len(runtime_addi) == len(request_ids) + ): + for i, rid in enumerate(request_ids): + if isinstance(rid, str) and isinstance(runtime_addi[i], dict): + runtime_addi_by_req[rid] = runtime_addi[i] + elif isinstance(request_ids, list) and isinstance(runtime_addi, dict): + for rid in request_ids: + if isinstance(rid, str) and isinstance(runtime_addi.get(rid), dict): + runtime_addi_by_req[rid] = runtime_addi[rid] + + # Containers to return per-request updates (e.g., code_predictor_hidden_per_request) + update_by_req_id: dict[str, dict] = {} + text_step = torch.zeros( + 1, + self.talker_config.text_config.hidden_size, + device=self._module_device(self.talker), + dtype=torch.bfloat16, + ) + last_talker_hidden = torch.zeros( + 1, + 1, + self.talker_config.text_config.hidden_size, + device=self._module_device(self.talker), + dtype=torch.bfloat16, + ) + + batched_talker_inputs = defaultdict(dict) + if is_profile: + batched_talker_inputs["0"] = { + "input_ids": input_ids, + "positions": positions, + "inputs_embeds": inputs_embeds.clone(), + "text_step": text_step, + "last_talker_hidden": last_talker_hidden, + } + + # ------- Prefill: span_len > 1 ------- + if ( + not is_profile + and isinstance(request_ids, list) + and isinstance(request_token_spans, list) + and isinstance(addi_by_req, dict) + ): + for idx_req, rid in enumerate(request_ids): + s, e = request_token_spans[idx_req] + span_len = int(e) - int(s) + if span_len <= 1: + continue + batched_talker_inputs[rid] = { + "input_ids": input_ids[s:e], + "positions": positions[s:e] if len(positions.shape) == 1 else positions[0, s:e], + "inputs_embeds": inputs_embeds[s:e].clone(), + "text_step": text_step, + "last_talker_hidden": last_talker_hidden, + } + info = addi_by_req.get(rid, {}) if isinstance(rid, str) else {} + # Read thinker outputs for prefill + thinker_sequence_embeds = info.get("thinker_embeddings").to( + device=self._module_device(self.talker), dtype=torch.bfloat16 + ) # Tensor [P,H] + thinker_hidden_states = info.get("thinker_hidden_states").to( + device=self._module_device(self.talker), dtype=torch.bfloat16 + ) # Tensor [K,H] + thinker_sequences = ( + info.get("thinker_sequences") + if info.get("thinker_sequences") is None + else torch.as_tensor(info.get("thinker_sequences"), device=self._module_device(self.talker)) + ) + thinker_chatml_ids = ( + info.get("thinker_input_ids") + if info.get("thinker_input_ids") is None + else torch.as_tensor(info.get("thinker_input_ids"), device=self._module_device(self.talker)) + ) + # Optional thinker-side TTS token embeddings (to be projected by talker) + tts_bos_thinker = info.get("tts_bos_embed").to( + device=self._module_device(self.talker), dtype=torch.bfloat16 + ) + tts_eos_thinker = info.get("tts_eos_embed").to( + device=self._module_device(self.talker), dtype=torch.bfloat16 + ) + tts_pad_thinker = info.get("tts_pad_embed").to( + device=self._module_device(self.talker), dtype=torch.bfloat16 + ) + + if thinker_sequence_embeds is None or thinker_hidden_states is None: + raise ValueError( + f"additional_information_by_req_id[{rid}] must include " + f"'thinker_embeddings' and 'thinker_hidden_states' for talker prefill." + ) + + # Normalize to tensors + if not isinstance(thinker_sequence_embeds, torch.Tensor): + thinker_sequence_embeds = torch.as_tensor( + thinker_sequence_embeds, device=self._module_device(self.talker) + ) + if not isinstance(thinker_hidden_states, torch.Tensor): + thinker_hidden_states = torch.as_tensor( + thinker_hidden_states, device=self._module_device(self.talker) + ) + + if isinstance(thinker_chatml_ids, torch.Tensor) or isinstance(thinker_chatml_ids, list): + ids_chatml = ( + thinker_chatml_ids + if isinstance(thinker_chatml_ids, torch.Tensor) + else torch.as_tensor(thinker_chatml_ids, device=self._module_device(self.talker)) + ) + if ids_chatml.ndim == 1: + ids_chatml = ids_chatml.unsqueeze(0) + else: + # Fallback: create dummy ids if not provided + ids_chatml = torch.zeros( + (1, thinker_sequence_embeds.shape[1]), + dtype=torch.long, + device=self._module_device(self.talker), + ) + thinker_sequences = ids_chatml + + speaker_id = self._get_text_spk_token_id(voice_type) + req_input_ids, req_embeds, trailing_text_hidden = self._thinker_to_talker_prefill( + thinker_embed=thinker_sequence_embeds.to(self._module_device(self.talker)), + thinker_hidden=thinker_hidden_states.to(self._module_device(self.talker)), + multimodal_mask=None, + input_ids=ids_chatml.to(self._module_device(self.talker)), + thinker_result_ids=thinker_sequences.to(self._module_device(self.talker)), + speaker_id=speaker_id, + tts_bos_thinker=tts_bos_thinker, + tts_eos_thinker=tts_eos_thinker, + tts_pad_thinker=tts_pad_thinker, + ) + # req_embeds: [S,D], req_input_ids: [S] → map到当前 span + seg_len = min(span_len, req_embeds.shape[0]) + # 强制对齐 dtype/device/内存布局，避免精度或布局差异导致比较不相等 + inputs_embeds[s : s + seg_len] = ( + req_embeds[:seg_len].to(device=inputs_embeds.device, dtype=inputs_embeds.dtype).contiguous() + ) + + if isinstance(req_input_ids, torch.Tensor) and req_input_ids.shape[0] >= seg_len: + input_ids[s : s + seg_len] = req_input_ids[:seg_len].to(input_ids.device) + # Queue tailing_text_hidden for decode (drop first for next steps), + # similar to Qwen2.5's thinker_reply_part_per_request + batched_talker_inputs[rid]["inputs_embeds"] = inputs_embeds[s : s + seg_len].clone() + try: + if isinstance(trailing_text_hidden, torch.Tensor) and trailing_text_hidden.numel() > 0: + if trailing_text_hidden.ndim == 2: + rem_tail = trailing_text_hidden + elif trailing_text_hidden.ndim == 1: + rem_tail = torch.zeros( + 0, + trailing_text_hidden.shape[0], + dtype=trailing_text_hidden.dtype, + device=trailing_text_hidden.device, + ) + else: + # compatible with old shape [1,S,D] + rem_tail = trailing_text_hidden.squeeze(0) + if rem_tail.shape[0] > 0: + update_by_req_id.setdefault(rid, {})["tailing_text_hidden_per_request"] = ( + rem_tail.detach().to("cpu").contiguous() + ) + # Also persist projected tts_pad for decode fallback if needed + if isinstance(tts_pad_thinker, torch.Tensor): + pad_in = tts_pad_thinker + if pad_in.ndim == 2: + pad_in = pad_in.unsqueeze(0) + if pad_in.ndim == 1: + pad_in = pad_in.view(1, 1, -1) + pad_proj = self.talker.text_projection(pad_in.to(self._module_device(self.talker))) + update_by_req_id.setdefault(rid, {})["tts_pad_embed_projected"] = ( + pad_proj.detach().to("cpu").contiguous() + ) + except Exception: + pass + + # ------- Decode: span_len == 1 ------- + if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list): + for idx_req, rid in enumerate(request_ids): + s, e = request_token_spans[idx_req] + if (int(e) - int(s)) != 1: + continue + batched_talker_inputs[rid] = { + "input_ids": input_ids[s:e], + "positions": positions[s:e] if len(positions.shape) == 1 else positions[0, s:e], + "inputs_embeds": inputs_embeds[s:e].clone(), + "text_step": text_step, + "last_talker_hidden": last_talker_hidden, + } + try: + q_tail = None + if isinstance(rid, str): + q_tail = runtime_addi_by_req.get(rid, {}).get("tailing_text_hidden_per_request") + if isinstance(q_tail, torch.Tensor) and q_tail.numel() > 0: + use_vec = q_tail[0:1, :] + new_q_tail = ( + q_tail[1:, :].detach().to("cpu").contiguous() + if q_tail.shape[1] > 1 + else self.tts_pad_embed.to(inputs_embeds.device, dtype=inputs_embeds.dtype) + ) + text_step = use_vec.to(inputs_embeds.device, dtype=inputs_embeds.dtype) + update_by_req_id.setdefault(rid, {})["tailing_text_hidden_per_request"] = new_q_tail + else: + text_step = self.tts_pad_embed.to(inputs_embeds.device, dtype=inputs_embeds.dtype) + + if isinstance(rid, str): + last_talker_hidden = ( + runtime_addi_by_req.get(rid, {}) + .get("last_talker_hidden") + .to(inputs_embeds.device, dtype=inputs_embeds.dtype) + ) + last_talker_hidden = last_talker_hidden.reshape( + *last_talker_hidden.shape[-2:] + ) # [1, hidden_size] + batched_talker_inputs[rid]["last_talker_hidden"] = last_talker_hidden + batched_talker_inputs[rid]["text_step"] = text_step + except Exception as e: + logger.error(f"Error in decode: {e}") + + if not is_profile and input_ids.shape[0] > 1: + self.input_ids_list = input_ids.tolist() + + last_talker_hidden = None + # Run talker forward + with torch.inference_mode(): + talker_hidden, batched_code_predictor_codes = self.talker( + batched_talker_inputs=batched_talker_inputs, + ) + if not is_profile: + for rid in request_ids: + idx_req = request_ids.index(rid) + s, e = request_token_spans[idx_req] + span_len = int(e) - int(s) + update_by_req_id.setdefault(rid, {})["last_talker_hidden"] = ( + talker_hidden[e - 1 : e, :].detach().to("cpu").contiguous() + ) + + multimodal_outputs: dict = None + # Return updates if any + if update_by_req_id: + multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id} + if not is_profile and multimodal_outputs is not None and batched_code_predictor_codes is not None: + multimodal_outputs.update({"code_predictor_codes": batched_code_predictor_codes}) + elif not is_profile and multimodal_outputs is None and batched_code_predictor_codes is not None: + multimodal_outputs = {"code_predictor_codes": batched_code_predictor_codes} + + next_token_id = None + + return OmniOutput( + text_hidden_states=talker_hidden, + multimodal_outputs=multimodal_outputs, + next_token_id=next_token_id, + ) + + # ========== Stage 3: Code2Wav ========== + elif self.model_stage == "code2wav": + # Extract codec codes from input + codes = [] + if input_ids is not None: + # addi_by_req = kwargs.get("additional_information_by_req_id") + codes.append(input_ids.reshape(1, 16, -1)) + + else: + # for profile, we use max length from inputs_embeds + codes.append( + torch.zeros( + (1, 16, inputs_embeds.shape[1]), + dtype=torch.long, + device=inputs_embeds.device, + ) + ) + + # # Remove EOS token if present + # if code[-1] == TALKER_CODEC_EOS_TOKEN_ID: + # code = code[:-1] + + # Generate audio from codec codes + audio_tensors = [] + for code in codes: + audio_tensor = self.generate_audio(code, voice_type) + audio_tensors.append(audio_tensor) + if len(audio_tensors) > 1: + logger.warning( + "Batched input for code2wav is not supported yet, only the first audio tensor will be returned" + ) + + return OmniOutput( + text_hidden_states=None, + multimodal_outputs={"model_outputs": audio_tensors[0].reshape(1, -1)}, + ) + + # Fallback (shouldn't reach here) + return OmniOutput( + text_hidden_states=torch.zeros( + [inputs_embeds.shape[0], self.talker.config.hidden_size], + dtype=torch.bfloat16, + ).to(self._module_device(self.model)), + multimodal_outputs=None, + ) + + # ==================== Audio Generation ==================== + + def generate_audio(self, code: torch.Tensor, voice_type: str) -> torch.Tensor: + """ + Generate audio waveform from codec codes. + + Args: + code: [8, T] - 8-layer RVQ codec codes + voice_type: Voice type (not used in Qwen3, kept for compatibility) + + Returns: + audio_tensor: [1, waveform_len] - Audio waveform + """ + code2wav_dev = self._module_device(self.code2wav) + + # Convert to tensor if needed + if isinstance(code, torch.Tensor): + talker_codes = code.to(dtype=torch.long, device=code2wav_dev) + else: + talker_codes = torch.as_tensor(code, dtype=torch.long, device=code2wav_dev) + + # Ensure shape is [batch=1, 8, T] + if talker_codes.ndim == 2: + # [8, T] → [1, 8, T] + talker_codes = talker_codes.unsqueeze(0) + elif talker_codes.ndim == 1: + # [T] → assume single layer, expand to 16 layers + talker_codes = talker_codes.unsqueeze(0).unsqueeze(0) + talker_codes = talker_codes.expand(1, 16, -1) + + # Use chunked decode for memory efficiency + audio_tensor = self.code2wav.chunked_decode( + talker_codes, + chunk_size=300, + left_context_size=25, + ) + + return audio_tensor + + # ==================== Thinker-Talker Projection ==================== + + def _load_talker_embedding(self) -> torch.nn.Embedding: + """Load talker embedding layer.""" + return self.talker.language_model.model.codec_embedding + + def _init_special_tokens_embeddings(self) -> set[str]: + """ + Initialize special token embeddings for thinker-talker projection. + + Following Transformers implementation: + - TTS tokens (BOS/EOS/PAD) come from thinker's embedding, projected to talker space + - Codec tokens (BOS/EOS/PAD/NOTHINK/THINK_*) come from talker's embedding + - Speaker tokens are also from talker's embedding + + Note on projections: + - text_projection: Used here for text token embeddings (thinker → talker dimension) + - hidden_projection: Used at runtime for multimodal hidden states (audio/image/video) + from thinker's last layer, not needed for special token initialization + """ + # Get embeddings from both models + # self.thinker_embedding = self.thinker.model.get_input_embeddings() + self.talker_embedding = self._load_talker_embedding() + + # Get configuration + talker_hf_config = self.talker_config + if hasattr(talker_hf_config, "talker_config"): + talker_hf_config = talker_hf_config.talker_config + + codec_special_tokens = torch.tensor( + [ + [ + talker_hf_config.codec_nothing_id, + talker_hf_config.codec_think_bos_id, + talker_hf_config.codec_think_eos_id, + talker_hf_config.codec_pad_id, + talker_hf_config.codec_bos_id, + talker_hf_config.codec_eos_token_id, + ] + ], + device=self._module_device(self.talker), + dtype=torch.long, + ) + codec_embeds = self.talker_embedding(codec_special_tokens) # [1, 6, talker_hidden] + ( + self.embed_codec_nothing_token, + self.embed_codec_think_bos_token, + self.embed_codec_think_eos_token, + self.embed_codec_pad_token, + self.embed_codec_bos_token, + self.embed_codec_eos_token, + ) = codec_embeds.chunk(6, dim=1) + + # Speaker token IDs (for voice selection) + # In Qwen3, speaker_id mapping is in talker_config.speaker_id + if hasattr(talker_hf_config, "speaker_id") and talker_hf_config.speaker_id: + self.tts_text_spk_token_ids = talker_hf_config.speaker_id + else: + # Default to audio_start_token_id if no speaker mapping + self.tts_text_spk_token_ids = { + "default": talker_hf_config.audio_start_token_id, + "Ethan": talker_hf_config.audio_start_token_id, + "prefix_caching": talker_hf_config.audio_start_token_id, + } + + self.default_tts_text_spk_type = list(self.tts_text_spk_token_ids.keys())[0] + + return set(["thinker_embedding.weight", "talker_embedding.weight"]) + + def _get_text_spk_token_id(self, voice_type: str) -> int: + """Get speaker token ID for voice type.""" + if voice_type not in self.tts_text_spk_token_ids: + return self.tts_text_spk_token_ids[self.default_tts_text_spk_type] + return self.tts_text_spk_token_ids[voice_type] + + def _thinker_to_talker_prefill( + self, + thinker_embed: torch.Tensor, + thinker_hidden: torch.Tensor, + multimodal_mask: Optional[torch.Tensor], + input_ids: torch.Tensor, + thinker_result_ids: torch.Tensor, + speaker_id, + tts_bos_thinker: Optional[torch.Tensor] = None, + tts_eos_thinker: Optional[torch.Tensor] = None, + tts_pad_thinker: Optional[torch.Tensor] = None, + ) -> tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]: + """ + Project thinker outputs to talker inputs during prefill stage. + + Returns: + (input_ids, input_embeds) for talker + """ + 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), + ), + dim=-1, + ) # Shape [n_starts + 1]; Take batch 0 since batched inference is not supported here. + multimodal_mask = ( + (thinker_result_ids == self.thinker_config.audio_token_id) | + (thinker_result_ids == self.thinker_config.image_token_id) | + (thinker_result_ids == self.thinker_config.video_token_id) + ).to(input_ids.device) # [t] # fmt: skip + + # Project thinker-side TTS embeddings into talker text space (provided by thinker stage) + talker_dev = self._module_device(self.talker) + + def _ensure_1x1(x: torch.Tensor) -> torch.Tensor: + if x.ndim == 3: + return x[0, -1:, :] + if x.ndim == 2: + return x[-1] + return x.view(1, 1, -1) + + def _proj_from_thinker(x_opt: Optional[torch.Tensor]) -> torch.Tensor: + if isinstance(x_opt, torch.Tensor) and x_opt.numel() > 0: + xin = _ensure_1x1(x_opt).to(talker_dev) + else: + xin = torch.zeros( + (1, thinker_embed.shape[-1]), + device=talker_dev, + dtype=thinker_embed.dtype, + ) + return self.talker.text_projection(xin).to(input_ids.device) + + tts_bos_embed = _proj_from_thinker(tts_bos_thinker) + tts_eos_embed = _proj_from_thinker(tts_eos_thinker) + tts_pad_embed = _proj_from_thinker(tts_pad_thinker) + self.tts_pad_embed = tts_pad_embed + + talker_input_embeds = [] # [1 t d] + talker_input_ids = [] + trailing_text_hidden_all: Optional[torch.Tensor] = None + # For every chatml parts + for i in range(len(im_start_indexes) - 1): + im_start_index = im_start_indexes[i].item() + segment_end_index = im_start_indexes[i + 1].item() + role_token = input_ids[0][im_start_index + 1] + # Talker should ignore thinker system prompt + if (role_token == self.config.system_token_id).item(): + continue + # Talker takes word embeddings for tokens and hidden state from `accept_hidden_layer` for multimodal inputs + elif (role_token == self.config.user_token_id).item(): + talker_user_part = self._get_talker_user_parts( + im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed + ) + talker_input_embeds.append(talker_user_part) + talker_input_ids.append(thinker_result_ids[im_start_index:segment_end_index]) + # Take assistant output (for now) + elif (role_token == self.config.assistant_token_id).item() and i == len(im_start_indexes) - 2: + talker_assistant_embeds, talker_assistant_ids, trailing_text_hidden = self._get_talker_assistant_parts( + im_start_index, + segment_end_index, + speaker_id, + thinker_embed, + tts_pad_embed, + tts_bos_embed, + tts_eos_embed, + ) + talker_input_embeds.append(talker_assistant_embeds) + talker_input_ids.append(talker_assistant_ids) + # capture trailing text hidden for decode steps + try: + if isinstance(trailing_text_hidden, torch.Tensor): + trailing_text_hidden_all = trailing_text_hidden + except Exception: + pass + # History assistant output (ignore for now) + elif (role_token == self.config.assistant_token_id).item() and i != len(im_start_indexes) - 2: + continue + else: + raise AssertionError("Expect role id after <|im_start|> (assistant, user, system)") + talker_input_embed = torch.cat([embed.to(input_ids.device) for embed in talker_input_embeds], dim=0) + talker_input_id = torch.cat([embed.to(input_ids.device) for embed in talker_input_ids], dim=0) + + return talker_input_id, talker_input_embed, trailing_text_hidden_all + + def _get_talker_user_parts(self, im_start_index, segment_end_index, multimodal_mask, thinker_hidden, thinker_embed): + user_talker_part = torch.empty( + (segment_end_index - im_start_index, self.config.talker_config.text_config.hidden_size), + device=thinker_hidden.device, + dtype=torch.bfloat16, + ) + + user_mm_mask = multimodal_mask[im_start_index:segment_end_index] + + # Multimodal data exists + if user_mm_mask.any(): + user_thinker_hidden_mm = thinker_hidden[im_start_index:segment_end_index][user_mm_mask] + mm_hidden = self.talker.hidden_projection(user_thinker_hidden_mm).to(thinker_hidden.device) + user_talker_part[user_mm_mask] = mm_hidden + user_thinker_embed = thinker_embed[im_start_index:segment_end_index][~user_mm_mask] + user_text_hidden = self.talker.text_projection(user_thinker_embed).to(thinker_hidden.device) + user_talker_part[~user_mm_mask] = user_text_hidden + return user_talker_part + + def _get_talker_assistant_parts( + self, im_start_index, segment_end_index, speaker_id, thinker_embed, tts_pad_embed, tts_bos_embed, tts_eos_embed + ): + assistant_hidden = self.talker.text_projection(thinker_embed[im_start_index:segment_end_index]).to( + tts_pad_embed.device + ) # [t, d] + assistant_text_hidden = torch.cat( + ( + assistant_hidden[:3], + tts_pad_embed.expand(4, -1), + tts_bos_embed, + assistant_hidden[3:4], # First text + ), + dim=0, + ) + codec_special_tokens = torch.tensor( + [ + self.config.talker_config.codec_nothing_id, + self.config.talker_config.codec_think_bos_id, + self.config.talker_config.codec_think_eos_id, + speaker_id, + self.config.talker_config.codec_pad_id, + self.config.talker_config.codec_bos_id, + ], + device=tts_pad_embed.device, + dtype=torch.long, + ) + assistant_codec_hidden = torch.cat( + ( + torch.zeros( + (3, self.config.talker_config.text_config.hidden_size), + device=tts_pad_embed.device, + dtype=torch.bfloat16, + ), + self.talker.get_input_embeddings(codec_special_tokens).to( + device=tts_pad_embed.device, dtype=torch.bfloat16 + ), + ), + dim=0, + ) + trailing_text_hidden = torch.cat( + ( + assistant_hidden[4:], + tts_eos_embed, + ), + dim=0, + ) + + input_embeds = assistant_text_hidden + assistant_codec_hidden + input_ids = torch.full( + (assistant_text_hidden.shape[0],), + fill_value=self.config.tts_pad_token_id, + dtype=torch.long, + device=assistant_text_hidden.device, + ) + return input_embeds, input_ids, trailing_text_hidden + + def _talker_to_code_predictor( + self, + talker_hidden_states: Optional[torch.Tensor], + layer0_token_ids: Optional[torch.Tensor], + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Project talker outputs to code predictor inputs. + + Returns: + (input_ids, input_embeds) for code predictor. + """ + predictor = getattr(self, "code_predictor", None) + device = ( + self._module_device(predictor) + if predictor is not None + else ( + talker_hidden_states.device + if isinstance(talker_hidden_states, torch.Tensor) + else torch.device("cuda" if torch.cuda.is_available() else "cpu") + ) + ) + + if not isinstance(talker_hidden_states, torch.Tensor): + raise ValueError("Talker hidden states must be provided for the code predictor stage.") + + inputs_embeds = talker_hidden_states.to(device=device, dtype=torch.bfloat16) + if inputs_embeds.ndim == 2: + inputs_embeds = inputs_embeds.unsqueeze(0) + + if not isinstance(layer0_token_ids, torch.Tensor): + raise ValueError("Layer-0 codec token ids must accompany talker hidden states.") + input_ids = layer0_token_ids.to(device=device, dtype=torch.long) + if input_ids.ndim == 1: + input_ids = input_ids.unsqueeze(0) + + return input_ids, inputs_embeds + + # ==================== Logits and Sampling ==================== + + def compute_logits( + self, + hidden_states: Union[torch.Tensor, OmniOutput], + sampling_metadata: SamplingMetadata = None, + ) -> Optional[torch.Tensor]: + """Compute logits from hidden states.""" + # Handle OmniOutput type + from vllm.v1.sample.logits_processor import LogitsProcessors + + if isinstance(hidden_states, OmniOutput): + hidden_states = hidden_states.text_hidden_states + num_reqs = hidden_states.shape[0] + top_k = hidden_states.size(1) - 1 + + def dummy_tensors(v): + return torch.full((num_reqs,), v, device=hidden_states.device) + + if sampling_metadata is None: + sampling_metadata = SamplingMetadata( + temperature=dummy_tensors(0.5), + all_greedy=False, + all_random=False, + top_p=dummy_tensors(0.9), + top_k=dummy_tensors(top_k), + generators={}, + max_num_logprobs=None, + no_penalties=True, + prompt_token_ids=None, + frequency_penalties=dummy_tensors(0.1), + presence_penalties=dummy_tensors(0.1), + repetition_penalties=dummy_tensors(0.1), + output_token_ids=[[] for _ in range(num_reqs)], + allowed_token_ids_mask=None, + bad_words_token_ids={}, + logitsprocs=LogitsProcessors(), + ) + # Use active model for logits computation + logits = self.model.compute_logits(hidden_states, sampling_metadata) # V, d + # Talker: suppress tokens by setting their probability to ~1e-9 (finite very small), + # implemented by assigning their logits to log(1e-9). + + if getattr(self, "model_stage", None) == "talker" and isinstance(logits, torch.Tensor): + # suppress tokens by setting their probability to ~1e-9 (finite very small) + suppressed_tokens = self._get_talker_suppressed_tokens() + try: + logits_cpu = logits.cpu() + logits_cpu[:, suppressed_tokens] = -1e9 + logits = logits_cpu.to(logits.device) + except Exception as e: + print(f"Error in logits suppression: {e}") + print(f"logits.shape: {logits.shape}") + print(f"suppressed_tokens: {suppressed_tokens}") + raise e + logits[:, suppressed_tokens] = -1e9 + return logits + + def sample( + self, + logits: torch.Tensor, + sampling_metadata: SamplingMetadata, + ) -> Optional[SamplerOutput]: + """Sample from logits.""" + return self.model.sample(logits, sampling_metadata) + + # ==================== Weight Loading ==================== + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights for all components of the omni model.""" + loaded_weights = set() + thinker_weights = [] + talker_weights = [] + code2wav_weights = [] + + # Separate weights by component + for k, v in weights: + if k.startswith("thinker."): + thinker_weights.append((k, v)) + elif k.startswith("talker."): + talker_weights.append((k, v)) + elif k.startswith("code2wav."): + code2wav_weights.append((k, v)) + else: + logger.warning(f"Unknown weight prefix: {k}") + # Load thinker weights + if self.thinker and thinker_weights: + thinker_loaded = self.thinker.load_weights(thinker_weights) + thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, "thinker") + loaded_weights.update(thinker_loaded) + + # Load talker weights + if self.talker and talker_weights: + talker_loaded = self.talker.load_weights(talker_weights) + talker_loaded = add_prefix_to_loaded_weights(talker_loaded, "talker") + loaded_weights.update(talker_loaded) + loaded_weights.update(self._init_special_tokens_embeddings()) + + # Load code2wav weights + if self.code2wav and code2wav_weights: + code2wav_loaded = self.code2wav.load_weights(code2wav_weights) + code2wav_loaded = add_prefix_to_loaded_weights(code2wav_loaded, "code2wav") + loaded_weights.update(code2wav_loaded) + + # Log summary + logger.info( + "Loaded %d weights for Qwen3OmniMoe (stage=%s)", + len(loaded_weights), + self.model_stage, + ) + + return loaded_weights diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py new file mode 100644 index 000000000..869d50ae5 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_code2wav.py @@ -0,0 +1,224 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2025 The Qwen team. +"""Inference-only Qwen3-Omni-Moe Code2Wav model.""" + +from __future__ import annotations + +from collections.abc import Iterable + +import numpy as np +import torch +import torch.nn as nn +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeCode2WavConfig, +) +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + Qwen3OmniMoeCausalConvNet, + Qwen3OmniMoeCausalTransConvNet, + Qwen3OmniMoeCode2WavDecoderBlock, + Qwen3OmniMoeCode2WavTransformerModel, + Qwen3OmniMoeConvNeXtBlock, + SnakeBeta, +) +from vllm.config import VllmConfig # type: ignore +from vllm.logger import init_logger # type: ignore +from vllm.model_executor.models.utils import ( # type: ignore + AutoWeightsLoader, + WeightsMapper, +) + +logger = init_logger(__name__) + + +class Qwen3OmniMoeCode2Wav(nn.Module): + """ + Qwen3 Omni MoE Code2Wav - Converts num_quantizers-layer RVQ codec codes to audio waveform. + + Architecture: + 1. Code Embedding: Embed and average num_quantizers RVQ layers + 2. Pre-Transformer: Add temporal context via sliding-window attention + 3. Upsampling: Progressive upsampling with ConvNeXt blocks + 4. Decoder: Multi-stage upsampling + residual units → waveform + + Input: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes + Output: [batch, 1, waveform_len] - Audio waveform [-1, 1] + + Total upsampling factor: ~1280x + Example: 100 codec frames → 128,000 audio samples (8 seconds at 16kHz) + """ + + input_modalities = "audio" + + # Weight mapper + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "code2wav.pre_transformer.": "pre_transformer.", + "code2wav.code_embedding.": "code_embedding.", + "code2wav.upsample.": "upsample.", + "code2wav.decoder.": "decoder.", + "code2wav.": "", + } + ) + + def __init__( + self, + *, + vllm_config: VllmConfig | None = None, + prefix: str = "", + ): + super().__init__() + + self.config: Qwen3OmniMoeCode2WavConfig = vllm_config.model_config.hf_config + + # Calculate total upsampling factor + self.total_upsample = np.prod(self.config.upsample_rates + self.config.upsampling_ratios) + + # Pre-transformer + self.pre_transformer = Qwen3OmniMoeCode2WavTransformerModel._from_config(self.config) + + # Code embedding: Single embedding table for all RVQ layers + self.code_embedding = nn.Embedding( + self.config.codebook_size * self.config.num_quantizers, self.config.hidden_size + ) + + # Offset for each RVQ layer (layer 0: 0-1023, layer 1: 1024-2047, etc.) + self.register_buffer( + "code_offset", + torch.arange(self.config.num_quantizers).view(1, -1, 1) * self.config.codebook_size, + persistent=False, + ) + + # Upsampling blocks (e.g., 2x, 2x) + upsample = [] + for factor in self.config.upsampling_ratios: + upsample.append( + nn.ModuleList( + [ + Qwen3OmniMoeCausalTransConvNet( + self.config.hidden_size, self.config.hidden_size, factor, factor + ), + Qwen3OmniMoeConvNeXtBlock(self.config.hidden_size), + ] + ) + ) + self.upsample = nn.ModuleList(upsample) + + # Decoder: Initial projection + progressive upsampling blocks + decoder = [Qwen3OmniMoeCausalConvNet(self.config.hidden_size, self.config.decoder_dim, kernel_size=7)] + + # Add decoder blocks (each upsamples and reduces channels) + for i in range(len(self.config.upsample_rates)): + decoder.append(Qwen3OmniMoeCode2WavDecoderBlock(self.config, i)) + + # Final projection to waveform + output_dim = self.config.decoder_dim // 2 ** len(self.config.upsample_rates) + decoder += [ + SnakeBeta(output_dim), + Qwen3OmniMoeCausalConvNet(output_dim, 1, kernel_size=7), + ] + self.decoder = nn.ModuleList(decoder) + + def forward(self, codes: torch.Tensor) -> torch.Tensor: + """ + Convert num_quantizers-layer RVQ codes to audio waveform. + + Args: + codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codec codes + + Returns: + waveform: [batch, 1, waveform_len] - Audio waveform clipped to [-1, 1] + """ + if codes.shape[1] != self.config.num_quantizers: + raise ValueError(f"Expected {self.config.num_quantizers} layers of codes, got {codes.shape[1]}") + + # Stage 1: Code Embedding + # Add offset to separate layer vocabularies, then embed and average + hidden = self.code_embedding(codes + self.code_offset).mean(1) + # Shape: [batch, seq_len, hidden_size] + + # Stage 2: Pre-Transformer (add temporal context) + hidden = self.pre_transformer(inputs_embeds=hidden).last_hidden_state + # Shape: [batch, seq_len, hidden_size] + + # Stage 3: Upsampling + hidden = hidden.permute(0, 2, 1) # [batch, hidden_size, seq_len] + for blocks in self.upsample: + for block in blocks: + hidden = block(hidden) + # Shape: [batch, hidden_size, seq_len * upsample_factor] + + # Stage 4: Decoder (progressive upsampling to waveform) + wav = hidden + for block in self.decoder: + wav = block(wav) + # Shape: [batch, 1, waveform_len] + + # Clamp to valid audio range + return wav.clamp(min=-1.0, max=1.0) + + def chunked_decode( + self, + codes: torch.Tensor, + chunk_size: int = 300, + left_context_size: int = 25, + ) -> torch.Tensor: + """ + Decode long sequences in chunks to avoid OOM. + + Uses overlapping chunks with left context to avoid boundary artifacts. + + Args: + codes: [batch, num_quantizers, seq_len] - num_quantizers-layer RVQ codes + chunk_size: Number of codec frames per chunk + left_context_size: Number of overlapping frames for context + + Returns: + waveform: [batch, 1, waveform_len] - Complete waveform + """ + wavs = [] + start_index = 0 + + while start_index < codes.shape[-1]: + end_index = min(start_index + chunk_size, codes.shape[-1]) + context_size = left_context_size if start_index >= left_context_size else start_index + + # Extract chunk with left context + codes_chunk = codes[..., start_index - context_size : end_index] + + # Decode chunk + wav_chunk = self(codes_chunk) + + # Remove context from output (context_size * total_upsample samples) + wavs.append(wav_chunk[..., context_size * self.total_upsample :]) + + start_index = end_index + + return torch.cat(wavs, dim=-1) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights from HuggingFace checkpoint.""" + loader = AutoWeightsLoader( + self, + skip_prefixes=["thinker.", "talker."], # Already loaded above + ) + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + # Log load summary + try: + total_bytes = 0 + for name, param in self.named_parameters(): + if param is not None and param.data is not None: + total_bytes += param.data.numel() * param.data.element_size() + device = next(self.parameters()).device + logger.info( + "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s", + self.__class__.__name__, + True, + total_bytes / (1024**2), + str(device), + ) + except Exception: + pass + + 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 new file mode 100644 index 000000000..49215be36 --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_code_predictor_mtp.py @@ -0,0 +1,491 @@ +"""Qwen3-Omni Code Predictor with MTP (Multi-Token Prediction) support. + +This module implements the code predictor component for Qwen3-Omni talker models. + +The code predictor generates residual RVQ (Residual Vector Quantization) codes +autoregressively, predicting layers 1 to N based on layer-0 codes from the talker. +""" + +from typing import Any, Optional + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers import Cache, PretrainedConfig +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeRotaryEmbedding +from vllm.config import CacheConfig, ModelConfig, VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.layernorm import RMSNorm +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding + +logger = init_logger(__name__) + + +# ============================================================================ +# Rotary Embeddings and Helper Functions +# ============================================================================ + + +class Qwen3OmniCodePredictorRotaryEmbedding(nn.Module): + """Rotary positional embeddings for the code predictor.""" + + def __init__(self, dim, max_position_embeddings=2048, base=10000): + super().__init__() + self.dim = dim + self.max_position_embeddings = max_position_embeddings + self.base = base + + inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2, dtype=torch.int64).float() / self.dim)) + self.register_buffer("inv_freq", inv_freq, persistent=False) + + def forward(self, x: torch.Tensor, position_ids: torch.LongTensor) -> tuple[torch.Tensor, torch.Tensor]: + inv_freq = self.inv_freq[None, :, None].float().expand(position_ids.shape[0], -1, 1) + position = position_ids[:, None, :].float() + freqs = (inv_freq @ position).transpose(1, 2) + emb = torch.cat((freqs, freqs), dim=-1) + cos = emb.cos().to(x.dtype) + sin = emb.sin().to(x.dtype) + return cos, sin + + +def rotate_half(x): + """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) + + +def apply_rotary_pos_emb(q, k, cos, sin, position_ids=None, unsqueeze_dim=1): + """Applies Rotary Position Embedding to the query and key tensors. + + Args: + q (`torch.Tensor`): The query tensor. + k (`torch.Tensor`): The key tensor. + cos (`torch.Tensor`): The cosine part of the rotary embedding. + sin (`torch.Tensor`): The sine part of the rotary embedding. + position_ids (`torch.Tensor`, *optional*): + Deprecated and unused. + unsqueeze_dim (`int`, *optional*, defaults to 1): + The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and + sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note + that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and + k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes + cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have + the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2. + Returns: + `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding. + """ + cos = cos.unsqueeze(unsqueeze_dim) + sin = sin.unsqueeze(unsqueeze_dim) + q_embed = (q * cos) + (rotate_half(q) * sin) + k_embed = (k * cos) + (rotate_half(k) * sin) + return q_embed, k_embed + + +def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor: + """Repeat key/value heads for grouped query attention.""" + if n_rep == 1: + return hidden_states + bsz, num_kv_heads, seq_len, head_dim = hidden_states.shape + hidden_states = hidden_states[:, :, None, :, :].expand(bsz, num_kv_heads, n_rep, seq_len, head_dim) + return hidden_states.reshape(bsz, num_kv_heads * n_rep, seq_len, head_dim) + + +# ============================================================================ +# Code Predictor Attention Layer +# ============================================================================ + + +class Qwen3OmniCodePredictorAttention(nn.Module): + """Multi-head self-attention for code predictor with vLLM optimization.""" + + def __init__(self, config, layer_idx: int, vllm_config: VllmConfig = None): + super().__init__() + + self.num_heads = config.code_predictor_config.num_attention_heads + self.num_key_value_heads = config.code_predictor_config.num_key_value_heads + self.head_dim = getattr( + config.code_predictor_config, + "head_dim", + config.code_predictor_config.hidden_size // config.code_predictor_config.num_attention_heads, + ) + self.hidden_size = config.code_predictor_config.hidden_size + + if self.num_heads % self.num_key_value_heads != 0: + raise ValueError("num_attention_heads must be divisible by num_key_value_heads") + + self.num_key_value_groups = self.num_heads // self.num_key_value_heads + + # Projection layers + 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_key_value_heads * self.head_dim, bias=False) + self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False) + self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False) + + # Query/Key normalization + self.q_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps) + self.k_norm = RMSNorm(self.head_dim, eps=config.code_predictor_config.rms_norm_eps) + self.is_causal = True + self.config = config + + def forward( + self, + hidden_states: torch.Tensor, + causal_mask: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + use_cache: bool = False, + position_ids: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + bsz, seq_len, _ = hidden_states.shape + + # Project to Q, K, V + q = self.q_proj(hidden_states).view(bsz, seq_len, self.num_heads, self.head_dim) + k = self.k_proj(hidden_states).view(bsz, seq_len, self.num_key_value_heads, self.head_dim) + v = self.v_proj(hidden_states).view(bsz, seq_len, self.num_key_value_heads, self.head_dim) + + # Apply normalization + q = self.q_norm(q) + k = self.k_norm(k) + + # Apply RoPE + q_heads = q.transpose(1, 2).contiguous() + k_heads = k.transpose(1, 2).contiguous() + q_heads, k_heads = apply_rotary_pos_emb(q_heads, k_heads, cos, sin) + v_heads = v.transpose(1, 2).contiguous() + + if past_key_values is not None: + # sin and cos are specific to RoPE models; cache_position needed for the static cache + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + k_heads, v_heads = past_key_values.update(k_heads, v_heads, self.layer_idx, cache_kwargs) + + from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS + + attention_interface = ALL_ATTENTION_FUNCTIONS["flash_attention_2"] + attn_output, _ = attention_interface( + self, + q_heads, + k_heads, + v_heads, + None, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.head_dim**-0.5, + sliding_window=None, + use_cache=use_cache, + position_ids=position_ids, + output_hidden_states=True, + output_attentions=False, + ) + attn_output = attn_output.reshape(*(hidden_states.shape[:-1]), -1).contiguous() + + attn_output = self.o_proj(attn_output) + return attn_output + + +# ============================================================================ +# Code Predictor MLP Layer +# ============================================================================ + + +class Qwen3OmniCodePredictorMLP(nn.Module): + """Feed-forward network for code predictor.""" + + def __init__(self, config): + super().__init__() + self.gate_proj = nn.Linear( + config.code_predictor_config.hidden_size, config.code_predictor_config.intermediate_size, bias=False + ) + self.up_proj = nn.Linear( + config.code_predictor_config.hidden_size, config.code_predictor_config.intermediate_size, bias=False + ) + self.down_proj = nn.Linear( + config.code_predictor_config.intermediate_size, config.code_predictor_config.hidden_size, bias=False + ) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + gate = F.silu(self.gate_proj(hidden_states)) + up = self.up_proj(hidden_states) + return self.down_proj(gate * up) + + +# ============================================================================ +# MTP Layer (Multi-Token Prediction Layer) +# ============================================================================ + + +class Qwen3OmniCodePredictorMTPLayer(nn.Module): + """MTP layer for speculative decoding - predicts next residual code layer.""" + + def __init__( + self, + config: PretrainedConfig, + prefix: str, + model_config: ModelConfig, + layer_idx: int, + cache_config: Optional[CacheConfig] = None, + quant_config: Optional[QuantizationConfig] = None, + ) -> None: + super().__init__() + self.layer_idx = layer_idx + self.config = config + + # Qwen3OmniCodePredictorDecoderLayer + self.self_attn = Qwen3OmniCodePredictorAttention( + config, + layer_idx, + vllm_config=type( + "VllmConfig", + (), + {"cache_config": cache_config, "quant_config": quant_config, "model_config": model_config}, + )(), + ) + self.mlp = Qwen3OmniCodePredictorMLP(config) + self.input_layernorm = RMSNorm( + config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps + ) + self.post_attention_layernorm = RMSNorm( + config.code_predictor_config.hidden_size, eps=config.code_predictor_config.rms_norm_eps + ) + + def mtp_block( + self, + hidden_states: torch.Tensor, + causal_mask: torch.Tensor, + cos: torch.Tensor, + sin: torch.Tensor, + past_key_values: Optional[Cache] = None, + cache_position: Optional[torch.LongTensor] = None, + use_cache: bool = False, + position_ids: Optional[torch.LongTensor] = None, + ) -> torch.Tensor: + # Self-attention with residual + residual = hidden_states + hidden_states = self.input_layernorm(hidden_states) + hidden_states = self.self_attn( + hidden_states, causal_mask, cos, sin, past_key_values, cache_position, use_cache, position_ids + ) + hidden_states = residual + hidden_states + + # MLP with residual + 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 + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + previous_hidden_states: torch.Tensor, + inputs_embeds: Optional[torch.Tensor] = None, + spec_step_index: int = 0, + ) -> torch.Tensor: + assert inputs_embeds is not None, "inputs_embeds required for MTP" + + # Mask position 0 (not needed for MTP) + inputs_embeds[positions == 0] = 0 + + hidden_states = torch.cat([inputs_embeds, previous_hidden_states], dim=-1) + + # Get position info for RoPE + batch_size, seq_len, _ = hidden_states.shape + device = hidden_states.device + position_ids = torch.arange(seq_len, dtype=torch.long, device=device).unsqueeze(0).expand(batch_size, -1) + + # Get RoPE embeddings + head_dim = self.self_attn.head_dim + rotary_emb = Qwen3OmniCodePredictorRotaryEmbedding( + head_dim, max_position_embeddings=self.config.code_predictor_config.max_position_embeddings + ) + rotary_emb = Qwen3OmniMoeRotaryEmbedding(self.config) + cos, sin = rotary_emb(hidden_states, position_ids) + + # Create causal mask + causal_mask = ( + torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1) + .unsqueeze(0) + .unsqueeze(0) + ) + causal_mask = causal_mask.masked_fill(causal_mask, float("-inf")) + + # Forward through MTP block + hidden_states = self.mtp_block(hidden_states, causal_mask, cos, sin) + + return hidden_states + + +class Qwen3OmniCodePredictorBaseModel(nn.Module): + """ + Base model for code predictor - matches HF Qwen3OmniMoeTalkerCodePredictorModel structure. + + This is a simple transformer that processes inputs_embeds and outputs hidden states. + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + config = vllm_config.model_config.hf_config.code_predictor_config + + self.config = config + self.vocab_size = config.vocab_size + self.num_code_groups = config.num_code_groups + + # Codec embeddings (for layers 1-num_code_groups-1) + self.codec_embedding = nn.ModuleList( + [ + VocabParallelEmbedding( + config.vocab_size, + config.hidden_size, + ) + for _ in range(config.num_code_groups - 1) + ] + ) + + # Decoder layers + self.layers = nn.ModuleList( + [ + Qwen3OmniCodePredictorMTPLayer( + vllm_config.model_config.hf_config, + f"{prefix}.layers.{idx}", + model_config=vllm_config.model_config, + layer_idx=idx, + cache_config=vllm_config.cache_config, + quant_config=vllm_config.quant_config, + ) + for idx in range(config.num_hidden_layers) + ] + ) + + self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + # RoPE + self.rotary_emb = Qwen3OmniMoeRotaryEmbedding(config=config) + + def forward( + self, + inputs_embeds: torch.Tensor, + attention_mask: Optional[torch.Tensor] = None, + position_ids: Optional[torch.LongTensor] = None, + past_key_values: Optional[Any] = None, + use_cache: Optional[bool] = False, + cache_position: Optional[torch.LongTensor] = None, + **kwargs, + ): + """ + Forward pass matching HF structure. + + Args: + inputs_embeds: [batch, seq_len, hidden_size] + + Returns: + Object with .last_hidden_state attribute + """ + batch_size, seq_len, _ = inputs_embeds.shape + + # Create positions tensor if not provided + # positions must be [num_tokens] or [batch_size, seq_len] + if position_ids is None: + if cache_position is not None: + position_ids = cache_position # [num_tokens] + else: + position_ids = torch.arange(seq_len, device=inputs_embeds.device).unsqueeze(0) # [1, seq_len] + else: + position_ids = position_ids.flatten() # Ensure [num_tokens] + + # Extract cos/sin from rotary_emb cache + # The cos_sin_cache is [max_pos, rotary_dim * 2] + cos, sin = self.rotary_emb(inputs_embeds, position_ids) + + # Create causal mask + device = inputs_embeds.device + causal_mask = ( + torch.triu(torch.ones(seq_len, seq_len, dtype=torch.bool, device=device), diagonal=1) + .unsqueeze(0) + .unsqueeze(0) + ) + causal_mask = causal_mask.masked_fill(causal_mask, float("-inf")) + + # Forward through decoder layers + hidden_states = inputs_embeds + + for layer in self.layers: + hidden_states = layer.mtp_block( + hidden_states, causal_mask, cos, sin, past_key_values, cache_position, use_cache, position_ids + ) + + # Final norm + hidden_states = self.norm(hidden_states) + + # Return in HF-compatible format + from collections import namedtuple + + Output = namedtuple("Output", ["last_hidden_state", "past_key_values"]) + return Output(last_hidden_state=hidden_states, past_key_values=None) # [batch, num_code_groups-1, hidden_size] + + def get_input_embeddings(self): + """Return codec embeddings for HF compatibility.""" + return self.codec_embedding + + +class Qwen3OmniMoeTalkerCodePredictor(nn.Module): + """ + Code predictor wrapper matching HF structure. + + Structure: + - self.model: Qwen3OmniCodePredictorBaseModel (transformer) + - self.lm_head: ModuleList of output heads + """ + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + + talker_code_predictor_config = vllm_config.model_config.hf_config + self.quant_config = vllm_config.quant_config + self.prefix = prefix + + self.config = talker_code_predictor_config + self.vocab_size = self.config.code_predictor_config.vocab_size + self.num_code_groups = self.config.code_predictor_config.num_code_groups + + # Base transformer model (matches HF structure) + self.model = Qwen3OmniCodePredictorBaseModel(vllm_config=vllm_config, prefix=prefix) + + # Output heads for each residual layer (1-num_layers-1) + self.lm_head = nn.ModuleList( + [ + nn.Linear( + self.config.code_predictor_config.hidden_size, + self.config.code_predictor_config.vocab_size, + bias=False, + ) + for _ in range(self.num_code_groups - 1) + ] + ) + + def forward( + self, + inputs_embeds: torch.Tensor, # [batch, seq_len, hidden_size] + layer_idx: int, # Which layer to predict (0-num_layers-2 for layers 1-num_layers-1) + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass for code predictor. + + Args: + inputs_embeds: Input embeddings [batch_size, seq_len, hidden_size] + layer_idx: Which residual layer to predict (0-num_layers-2 for layers 1-num_layers-1) + + Returns: + logits: Predicted logits [batch_size, seq_len, vocab_size] + hidden_states: Output hidden states [batch_size, seq_len, hidden_size] + """ + # Pass through base model + hidden_states = self.model(inputs_embeds) + + # Get logits from corresponding head + logits = self.lm_head[layer_idx](hidden_states) + + return logits, hidden_states diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py new file mode 100644 index 000000000..34e05925f --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -0,0 +1,703 @@ +from collections.abc import Iterable + +import torch +import torch.nn as nn +import torch.nn.functional as F +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeTalkerConfig, +) +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + Qwen3OmniMoeAudioEncoder, +) +from vllm.config import VllmConfig +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY +from vllm.model_executor.layers.linear import ReplicatedLinear +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.qwen2_5_omni_thinker import ( + Qwen2_5OmniThinkerDummyInputsBuilder, +) +from vllm.model_executor.models.qwen3_moe import Qwen3MoeMLP +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + maybe_prefix, + sequence_parallel_chunk, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.sequence import IntermediateTensors +from vllm.v1.sample.metadata import SamplingMetadata + +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp import ( + Qwen3OmniMoeTalkerCodePredictor, +) +from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( + Qwen3MoeLLMForCausalLM, + Qwen3Omni_VisionTransformer, + Qwen3OmniMoeConditionalGenerationMixin, + Qwen3OmniMoeThinkerMultiModalProcessor, + Qwen3OmniMoeThinkerProcessingInfo, +) + +try: + import flash_attn +except (ImportError, ModuleNotFoundError): + flash_attn = None + +from vllm_omni.model_executor.models.utils import safe_tensor_reshape + +logger = init_logger(__name__) + +Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3OmniMoeThinkerMultiModalProcessor, + info=Qwen3OmniMoeThinkerProcessingInfo, + dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder, +) +class Qwen3OmniMoeTalkerForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + Qwen3OmniMoeConditionalGenerationMixin, +): + """ + Qwen3 Omni MoE Talker - Converts text to audio codec codes. + + The talker is the second stage of Qwen3 Omni MoE's TTS pipeline: + 1. Thinker: Generates text response + hidden states + 2. Talker: Converts those to 8-layer audio codec codes + 3. Code2Wav: Converts codes to waveform + + ## Key Components: + - text_projection: Projects thinker text embeddings → talker dimension + - hidden_projection: Projects thinker hidden states → talker dimension + - language_model: Main MoE transformer (generates layer 0) + - codec_head: Projects to codec vocabulary (layer 0 logits) + - code_predictor: Small transformer for layers 1-num_layers-1 + """ + + logger = init_logger(__name__) + + # Weight mapping from HuggingFace to vLLM naming convention + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + # Main MoE transformer model + "talker.model.": "language_model.model.", + # Codec head remains separate (outputs audio codes, not text) + "talker.codec_head.": "codec_head.", + # Code predictor: Now matches HF structure exactly (has .model sub-module) + # e.g., "talker.code_predictor.model.codec_embedding.0" → "code_predictor.model.codec_embedding.0" + "talker.code_predictor.": "code_predictor.", + # Projection layers + "talker.text_projection.": "text_projection.", + "talker.hidden_projection.": "hidden_projection.", + # Fallback: strip talker prefix + "talker.": "", + } + ) + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config + self.quant_config = vllm_config.quant_config + self.prefix = prefix + + self.config = talker_config + self.vocab_size = talker_config.text_config.vocab_size + self.router_aux_loss_coef = talker_config.text_config.router_aux_loss_coef + self.num_experts = talker_config.text_config.num_experts + self.num_experts_per_tok = talker_config.text_config.num_experts_per_tok + # thinker projection components for talker + self.text_projection = Qwen3OmniMoeTalkerResizeMLP(self.config) + self.hidden_projection = Qwen3OmniMoeTalkerResizeMLP(self.config) + self.codec_head = nn.Linear(self.config.text_config.hidden_size, self.config.text_config.vocab_size, bias=False) + + self.rope_deltas = None + self.spatial_merge_size = self.config.spatial_merge_size + + self.vocab_size = self.config.code_predictor_config.vocab_size + self.num_code_groups = self.config.code_predictor_config.num_code_groups + + self.language_model = Qwen3OmniMoeModel( + vllm_config=vllm_config, + talker_config=self.config, + prefix=maybe_prefix(prefix, "language_model"), + ) + + self.code_predictor = Qwen3OmniMoeTalkerCodePredictor( + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "code_predictor") + ) + + def code_predictor_forward( + self, + input_ids: torch.Tensor, + inputs_embeds: torch.Tensor | None = None, + *, + temperature: float = 1.0, + top_k: int = 50, # Match transformers default + top_p: float = 0.8, # Match transformers default + generation_steps: int | None = None, + last_talker_hidden: torch.Tensor | None = None, + **_: object, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Generate full RVQ codec codes for the provided sequence. + + The code predictor consumes the layer-0 codec codes produced by the talker + alongside the talker's hidden states, and autoregressively predicts the remaining + residual layers (to num_codec_groups). + + Returns: + tuple containing: + - residual_codes: A tensor of shape [batch, num_code_groups, seq_len] containing + the complete set of codec codes + - summed_embeddings: A tensor of shape [batch, seq_len, hidden_size] + Sum of all layer embeddings at each position (like Transformers) + """ + if input_ids is None: + raise ValueError("`input_ids` containing layer-0 codec codes must be provided.") + if inputs_embeds is None: + raise ValueError("`inputs_embeds` containing talker hidden states must be provided.") + + if inputs_embeds.ndim == 2: + inputs_embeds = inputs_embeds.unsqueeze(0) + if input_ids.ndim == 1: + input_ids = input_ids.unsqueeze(0) + + # Ensure the tensors are contiguous for the autoregressive sampling loop + inputs_embeds = inputs_embeds.contiguous() + input_ids = input_ids.contiguous() + + # Generate full codec codes using MTP + # This will be the parallel prediction implementation + batch_size, seq_len = input_ids.shape + + # For now, use sequential generation (TODO: implement parallel) + # Result will be [batch, num_code_groups, seq_len] + # - all_codes_per_position will collect [batch, num_code_groups, 1] for each position + all_codes_per_position = [] + middle_hidden_states = [] # Collect hidden states for each position + + # Generate residual layers for each position + for pos in range(seq_len): + layer0_code = input_ids[:, pos : pos + 1] # [batch, 1] + + # Predict all residual layers (layers 1 to num_code_groups-1) autoregressively + pos_codes = [layer0_code] # Start with layer 0: [batch, 1] + + # Initial input: [last_talker_hidden, layer0_embed] + layer0_embed = self.get_input_embeddings(layer0_code) + prev_embed = layer0_embed # Track previous layer embedding + try: + current_input = torch.cat([last_talker_hidden, prev_embed], dim=1) # [batch, 2, hidden_size] + except Exception as e: + print(f"Error in current_input: {e}") + print(f"last_talker_hidden shape: {last_talker_hidden.shape}") + print(f"prev_embed shape: {prev_embed.shape}") + raise e + + for layer_idx in range(self.num_code_groups - 1): + # Input for this layer: [last_talker_hidden, prev_layer_embed] + + # Forward through code_predictor model + outputs = self.code_predictor.model( + inputs_embeds=current_input, + attention_mask=None, + position_ids=None, + past_key_values=None, + use_cache=False, + cache_position=None, + ) + + hidden_state = outputs.last_hidden_state # [batch, 2, hidden_size] + + # Use the corresponding lm_head for this layer + logits = self.code_predictor.lm_head[layer_idx](hidden_state[:, -1:, :]) # [batch, 1, vocab_size] + from transformers.generation.logits_process import ( + LogitsProcessorList, + TopKLogitsWarper, + TopPLogitsWarper, + ) + + logits_processors = LogitsProcessorList( + [ + TopKLogitsWarper(top_k=top_k), + TopPLogitsWarper(top_p=top_p), + ] + ) + input_ids_for_logits_processors = torch.tensor([pos_codes[1:]]).to(logits.device, dtype=torch.long) + logits = logits_processors(input_ids_for_logits_processors, logits.squeeze(0)).unsqueeze(0) + + # Sample from the filtered distribution + probs = F.softmax(logits, dim=-1) + code = torch.multinomial(probs.squeeze(1), num_samples=1) # [batch, 1] + pos_codes.append(code) + + # Update prev_embed for next layer (if not last layer) + # layer_idx=0 predicts layer 1, embed with codec_embedding[1] + new_embed = self.code_predictor.model.codec_embedding[layer_idx](code) # [batch, 1, hidden_size] + current_input = torch.cat([current_input, new_embed], dim=1) # [batch, 3~n, hidden_size] + + # Stack all layers for this position: [batch, num_code_groups, 1] + pos_all_layers = torch.stack(pos_codes, dim=1) # [batch, num_code_groups, 1] + all_codes_per_position.append(pos_all_layers) + middle_hidden_states.append(current_input[:, 2:-1, :]) + + # Concatenate across positions: [batch, num_code_groups, seq_len] + result_codes = torch.cat(all_codes_per_position, dim=2) + + # Build summed embeddings for each position (like Transformers) + # This combines layer-0 embed, mid layers hidden states, and last layer embed + all_summed_embeddings = [] + + for pos in range(seq_len): + # Layer 0 embedding + layer0_code = result_codes[:, 0, pos : pos + 1] # [batch, 1] + layer0_embed = self.get_input_embeddings(layer0_code) # [batch, 1, hidden_size] + + # mid layers hidden states (from CodePredictor) + mid_residual_hiddens = middle_hidden_states[pos] # [batch, num_code_groups-2, hidden_size] + mid_list = list(mid_residual_hiddens.split(1, dim=1)) + + # last layer embedding + last_layer_code = result_codes[:, -1, pos : pos + 1] # [batch, 1] + last_residual_hidden = self.code_predictor.model.codec_embedding[-1](last_layer_code) + + # Concatenate all layers: [batch, num_code_groups, hidden_size] + pos_codec_hiddens = torch.cat( + [layer0_embed] + mid_list + [last_residual_hidden], + dim=1, + ) + + # Sum across layers: [batch, 1, hidden_size] (like Transformers) + pos_summed = pos_codec_hiddens.sum(dim=1, keepdim=True) + all_summed_embeddings.append(pos_summed) + + # Concatenate across positions: [batch, seq_len, hidden_size] + summed_embeddings = torch.cat(all_summed_embeddings, dim=1) + + return result_codes, summed_embeddings + + def init_multi_modal(self, thinker_config): + """ + Initialize multimodal components from the thinker. + + Unlike Qwen2.5 Omni which creates audio_tower and visual encoders here, + Qwen3 Omni MoE has a cleaner separation: the thinker is the ONLY module + that processes raw multimodal inputs. The talker only handles text-to-audio + conversion using pre-processed embeddings from the thinker. + + This method exists for API compatibility and stores the thinker config + for reference. The actual multimodal processing components (audio_tower, + visual) are ONLY in the thinker, not duplicated in the talker. + + Args: + thinker_config: Configuration from the thinker model (for reference only) + """ + self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config) + self.visual = Qwen3Omni_VisionTransformer( + vision_config=thinker_config.vision_config, + norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), + quant_config=self.quant_config, + prefix=maybe_prefix(self.prefix, "visual"), + # attn_backend_override=attn_backend_override, + ) + + def project_thinker_outputs( + self, + thinker_embeds: torch.Tensor | None = None, + thinker_hidden_states: torch.Tensor | None = None, + is_multimodal_mask: torch.Tensor | None = None, + ) -> torch.Tensor: + """ + Project thinker outputs to talker's hidden dimension. + + The talker has a different hidden size than the thinker, so we need + to project the inputs appropriately: + - Text embeddings (from thinker's embedding layer) → text_projection + - Hidden states (from thinker's last layer, for multimodal) → hidden_projection + + Args: + thinker_embeds: Text embeddings from thinker [batch, seq, thinker_hidden] + thinker_hidden_states: Hidden states from thinker's last layer [batch, seq, thinker_hidden] + is_multimodal_mask: Boolean mask indicating multimodal positions [batch, seq] + + Returns: + projected_embeds: [batch, seq, talker_hidden] + """ + if thinker_embeds is None and thinker_hidden_states is None: + raise ValueError("Either thinker_embeds or thinker_hidden_states must be provided") + + # If only embeddings provided, project all as text + if thinker_hidden_states is None or is_multimodal_mask is None: + return self.text_projection(thinker_embeds) + + # If only hidden states provided, project all as hidden + if thinker_embeds is None: + return self.hidden_projection(thinker_hidden_states) + + # Mixed case: use mask to decide which projection + batch_size, seq_len, _ = thinker_embeds.shape + output = torch.empty( + (batch_size, seq_len, self.config.text_config.hidden_size), + device=thinker_embeds.device, + dtype=thinker_embeds.dtype, + ) + + # Project multimodal regions using hidden states + if is_multimodal_mask.any(): + mm_hidden = thinker_hidden_states[is_multimodal_mask] + projected_mm = self.hidden_projection(mm_hidden) + output[is_multimodal_mask] = projected_mm + + # Project text regions using embeddings + if (~is_multimodal_mask).any(): + text_embeds = thinker_embeds[~is_multimodal_mask] + projected_text = self.text_projection(text_embeds) + output[~is_multimodal_mask] = projected_text + + return output + + def forward( + self, + batched_talker_inputs: dict, + intermediate_tensors: IntermediateTensors | None = None, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + """Forward pass through the talker model. + + For inference, the talker receives inputs_embeds that should already be + projected to talker's hidden dimension. If receiving raw thinker outputs, + use project_thinker_outputs() first. + """ + # If intermediate_tensors is provided (pipeline parallel), + # inputs_embeds should be None + batched_input_ids = [] + batched_positions = [] + batched_input_embeds = [] + batched_code_predictor_codes = [] + if intermediate_tensors is not None: + inputs_embeds = None + for rid in batched_talker_inputs: + input_ids = safe_tensor_reshape(batched_talker_inputs[rid]["input_ids"], (1, -1)) + positions = batched_talker_inputs[rid]["positions"] + inputs_embeds = safe_tensor_reshape( + batched_talker_inputs[rid]["inputs_embeds"], (-1, self.config.text_config.hidden_size) + ) + text_step = safe_tensor_reshape(batched_talker_inputs[rid]["text_step"], (1, -1)) + last_talker_hidden = safe_tensor_reshape( + batched_talker_inputs[rid]["last_talker_hidden"], (1, 1, self.config.text_config.hidden_size) + ) + # for profiling + if inputs_embeds.shape[-1] == 2048: + inputs_embeds = self.text_projection(inputs_embeds) + if inputs_embeds.shape[0] == 1: + code_predictor_codes, summed_embeddings = self.code_predictor_forward( + input_ids, inputs_embeds.clone(), last_talker_hidden=last_talker_hidden + ) + inputs_embeds = summed_embeddings.clone() + else: + code_predictor_codes = torch.zeros((0, self.num_code_groups), dtype=torch.long) + batched_input_ids.append(input_ids.reshape(-1)) + batched_positions.append(positions) + batched_input_embeds.append((inputs_embeds + text_step).reshape(-1, self.config.text_config.hidden_size)) + batched_code_predictor_codes.append(code_predictor_codes.squeeze(-1).detach().to("cpu").contiguous()) + try: + talker_input_ids = torch.cat(batched_input_ids, dim=0) + talker_positions = torch.cat(batched_positions, dim=0) + talker_input_embeds = torch.cat(batched_input_embeds, dim=0) + except Exception as e: + print(f"Error in talker_input_embeds: {e}") + print(f"talker_input_embeds shape: {[tensor.shape for tensor in batched_input_embeds]}") + print(f"talker_input_ids shape: {[tensor.shape for tensor in batched_input_ids]}") + print(f"talker_positions shape: {[tensor.shape for tensor in batched_positions]}") + raise e + talker_hidden_states, _ = self.language_model.model( + talker_input_ids, + talker_positions, + intermediate_tensors, + inputs_embeds=talker_input_embeds, + **kwargs, + ) + + # Pass talker hidden states to code predictor + # Returns: (residual_codes, summed_embeddings=sum of all layer embeddings at each position) + + # Return both talker hidden states and code predictor results + # code_predictor_codes: [batch, num_code_groups, seq_len] + # summed_embeddings: [batch, seq_len, hidden_size] + # - Sum of all layer embeddings at each position (like Transformers) + + return talker_hidden_states, batched_code_predictor_codes + + def compute_logits( + self, + hidden_states: torch.Tensor, + sampling_metadata: SamplingMetadata = None, + ) -> torch.Tensor | None: + """Compute logits for audio codec codes (layer 0 of RVQ). + + This projects the hidden states to the codec vocabulary space. + For full audio generation, layers except 0 would be predicted by + the code_predictor after sampling. + """ + logits = self.codec_head(hidden_states) + return logits + + def make_empty_intermediate_tensors( + self, batch_size: int, dtype: torch.dtype, device: torch.device + ) -> IntermediateTensors: + """Create empty intermediate tensors for pipeline parallelism.""" + return self.language_model.make_empty_intermediate_tensors(batch_size, dtype, device) + + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: + mm_input_by_modality = {} + + # Preserve the order of modalities if there are multiple of them + # from the order of kwargs. + for input_key in kwargs: + if input_key in ("pixel_values", "image_embeds") and "image" not in mm_input_by_modality: + mm_input_by_modality["image"] = self._parse_and_validate_image_input(**kwargs) + if input_key in ("pixel_values_videos", "video_embeds") and "video" not in mm_input_by_modality: + mm_input_by_modality["video"] = self._parse_and_validate_video_input(**kwargs) + if input_key in ("input_audio_features") and "audio" not in mm_input_by_modality: + mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs) + return mm_input_by_modality + + def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings | None: + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) + if not mm_input_by_modality: + return [] + + # The result multimodal_embeddings is tuple of tensors, with each + # tensor correspoending to a multimodal data item (image or video). + multimodal_embeddings: tuple[torch.Tensor, ...] = () + + # NOTE: It is important to iterate over the keys in this dictionary + # to preserve the order of the modalities. + # TODO: do projection for all multimodel + for modality in mm_input_by_modality: + multimodal_input = mm_input_by_modality[modality] + if modality == "image": + image_embeddings = self._process_image_input(multimodal_input) + image_embeddings = self.hidden_projection(image_embeddings) + multimodal_embeddings += tuple(image_embeddings) + if modality == "video": + video_embeddings = self._process_video_input(multimodal_input) + video_video_embeddings_project = () + for video_embed in video_embeddings: + proj = nn.Linear(8192, 2048).to(device=video_embed.device, dtype=torch.bfloat16) + video_embed = proj(video_embed) + video_embed_project = self.hidden_projection(video_embed) + video_video_embeddings_project += (video_embed_project,) + multimodal_embeddings += tuple(video_video_embeddings_project) + if modality == "audio": + audio_embeddings = self._process_audio_input(multimodal_input) + audio_embeddings = self.hidden_projection(audio_embeddings) + multimodal_embeddings += tuple(audio_embeddings) + return multimodal_embeddings + + def get_input_embeddings(self, input_ids: torch.Tensor, multimodal_embeddings: MultiModalEmbeddings | None = None): + """Get the input embedding layer (for codec tokens).""" + return self.language_model.get_input_embeddings(input_ids, multimodal_embeddings) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load weights for the talker model. + + The weight mapping translates from HuggingFace naming convention + to vLLM's internal structure. Code predictor weights are routed + to its custom loader for vocab extension support. + """ + loader = AutoWeightsLoader( + self, + skip_prefixes=["thinker.", "code2wav."], + # "code_predictor."], + ) + # Don't apply mapper again since we already did it + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + # Log load summary + try: + total_bytes = 0 + for name, param in self.named_parameters(): + if param is not None and param.data is not None: + total_bytes += param.data.numel() * param.data.element_size() + device = next(self.parameters()).device + logger.info( + "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s", + self.__class__.__name__, + True, + total_bytes / (1024**2), + str(device), + ) + except Exception: + pass + + multi_model_weights = set() + for name, param in self.visual.named_parameters(): + multi_model_weights.add("visual." + name) + for name, param in self.audio_tower.named_parameters(): + multi_model_weights.add("audio_tower." + name) + loaded.update(multi_model_weights) + + return loaded + + +class Qwen3OmniMoeTalkerResizeMLP(nn.Module): + """ + MLP for projecting between thinker and talker hidden dimensions. + + The thinker and talker have different hidden sizes: + - Thinker: config.thinker_hidden_size (e.g., 3584) + - Talker: config.text_config.hidden_size (e.g., 2048) + + This MLP projects from thinker → talker dimension. + Two instances are used: + - text_projection: For text embeddings from thinker's embedding layer + - hidden_projection: For hidden states from thinker's last transformer layer + """ + + def __init__(self, config: Qwen3OmniMoeTalkerConfig): + super().__init__() + self.linear_fc1 = nn.Linear(config.thinker_hidden_size, config.text_config.intermediate_size, bias=True) + self.linear_fc2 = nn.Linear(config.text_config.intermediate_size, config.text_config.hidden_size, bias=True) + self.act_fn = _ACTIVATION_REGISTRY[config.text_config.hidden_act] # silu + + def forward(self, hidden_state): + return self.linear_fc2(self.act_fn(self.linear_fc1(hidden_state))) + + +class Qwen3OmniMoeModel(Qwen3MoeLLMForCausalLM): + def __init__(self, vllm_config, talker_config, prefix): + talker_vllm_config = vllm_config.with_hf_config( + talker_config.text_config, architectures=["Qwen3MoeForCausalLM"] + ) + talker_vllm_config.model_config.hf_text_config = talker_vllm_config.model_config.hf_config + super().__init__( + vllm_config=talker_vllm_config, + prefix=prefix, + ) + + self.config = talker_config + + # Remove the inherited LM head so the talker only exposes codec outputs. + if hasattr(self, "lm_head"): + del self.lm_head + + # Replace the base embed tokens with codec embedding (defined below). + if hasattr(self.model, "embed_tokens"): + del self.model.embed_tokens + + # Codec embedding for RVQ code generation + self.model.codec_embedding = nn.Embedding( + talker_config.text_config.vocab_size, talker_config.text_config.hidden_size + ) + + # Add shared expert to each MoE layer and patch the forward method + layer_idx = 0 + for layer in self.model.layers: + # add shared expert to Qwen3OmniMoeSparseMoeBlock layers + if hasattr(layer.mlp, "experts"): # Check if it's a SparseMoeBlock + # Shared expert is a regular gated MLP (SwiGLU) + layer.mlp.shared_expert = Qwen3MoeMLP( + hidden_size=self.config.text_config.hidden_size, + intermediate_size=self.config.text_config.shared_expert_intermediate_size, + hidden_act=self.config.text_config.hidden_act, + quant_config=talker_vllm_config.quant_config, + reduce_results=False, # Don't reduce since we'll add it manually + prefix=f"{prefix}.layers.{layer_idx}.mlp.shared_expert", + ) + + # Shared expert gate outputs a single scalar per token + layer.mlp.shared_expert_gate = ReplicatedLinear( + self.config.text_config.hidden_size, + 1, # Output single scalar per token + bias=False, + quant_config=None, + prefix=f"{prefix}.layers.{layer_idx}.mlp.shared_expert_gate", + ) + + # Store MoE config values for router computation + layer.mlp.top_k = self.config.text_config.num_experts_per_tok + layer.mlp.norm_topk_prob = self.config.text_config.norm_topk_prob + layer.mlp.num_experts = self.config.text_config.num_experts + + # Monkey-patch the forward method to use shared expert + layer.mlp.forward = self._create_moe_forward_with_shared_expert(layer.mlp) + + layer_idx += 1 + + def _create_moe_forward_with_shared_expert(self, moe_layer): + """Create a forward method that includes shared expert computation. + + This matches the Transformers implementation where: + 1. Compute shared expert output (regular MLP) + 2. Gate it with sigmoid(shared_expert_gate(x)) + 3. Apply softmax BEFORE top-k selection (matches Transformers router) + 4. Add to routed expert outputs + """ + + def forward_with_shared_expert(hidden_states: torch.Tensor, layer_idx: int = 0) -> torch.Tensor: + # Save original shape + orig_shape = hidden_states.shape + hidden_dim = hidden_states.shape[-1] + hidden_states = hidden_states.view(-1, hidden_dim) + + # handle sequence parallel if needed + if hasattr(moe_layer, "is_sequence_parallel") and moe_layer.is_sequence_parallel: + hidden_states = sequence_parallel_chunk(hidden_states) + + # Compute shared expert output + # The shared expert is a regular MLP, not a routed MoE + shared_output = None + if hasattr(moe_layer, "shared_expert") and moe_layer.shared_expert is not None: + # Forward through shared expert MLP + shared_output = moe_layer.shared_expert(hidden_states) + + # Apply gating with sigmoid: sigmoid(gate(x)) * shared_expert(x) + if hasattr(moe_layer, "shared_expert_gate") and moe_layer.shared_expert_gate is not None: + gate_logits, _ = moe_layer.shared_expert_gate(hidden_states) + gate_values = F.sigmoid(gate_logits) # [batch, 1] + shared_output = gate_values * shared_output # Broadcasting: [batch, 1] * [batch, hidden] + + # Compute experts results + # router_logits: (num_tokens, n_experts) + router_logits, _ = moe_layer.gate(hidden_states) + experts_output = moe_layer.experts(hidden_states=hidden_states, router_logits=router_logits) + + # combine experts and shared expert results + if shared_output is not None: + final_hidden_states = experts_output + shared_output + + # Handle sequence parallel if needed + if hasattr(moe_layer, "is_sequence_parallel") and moe_layer.is_sequence_parallel: + from vllm.distributed import tensor_model_parallel_all_gather + + num_tokens = orig_shape[0] if len(orig_shape) > 1 else 1 + final_hidden_states = tensor_model_parallel_all_gather(final_hidden_states, 0) + final_hidden_states = final_hidden_states[:num_tokens] + try: + final_hidden_states.view(orig_shape) + except Exception as e: + print(f"Error viewing final hidden states: {e}") + print(f"final_hidden_states.shape: {final_hidden_states.shape}") + print(f"orig_shape: {orig_shape}") + raise e + # Return with original shape + return final_hidden_states.view(orig_shape) + + return forward_with_shared_expert + + def get_input_embeddings( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + generation_steps=None, + ) -> torch.Tensor: + return self.model.codec_embedding(input_ids) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py new file mode 100644 index 000000000..573dc42cc --- /dev/null +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -0,0 +1,1668 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2025 The Qwen team. +# Copyright 2023 The vLLM team. +# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved. +# +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX +# and OPT implementations in this library. It has been modified from its +# original forms to accommodate minor architectural differences compared +# to GPT-NeoX and OPT used by the Meta AI team that trained the model. +# +# 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. +"""Inference-only Qwen3-Omni-Moe model (thinker part).""" + +import math +from collections.abc import Callable, Iterable, Mapping, Sequence +from functools import partial +from typing import Any + +import numpy as np +import torch +import torch.nn as nn +import torch.nn.functional as F +from packaging.version import Version +from transformers import PretrainedConfig +from transformers import __version__ as TRANSFORMERS_VERSION +from transformers.feature_extraction_utils import BatchFeature +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( + Qwen3OmniMoeConfig, + Qwen3OmniMoeThinkerConfig, +) +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( + Qwen3OmniMoeAudioEncoder, +) +from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import ( + Qwen3OmniMoeProcessor, +) +from transformers.models.whisper import WhisperFeatureExtractor +from vllm.attention.utils.fa_utils import get_flash_attn_version +from vllm.compilation.decorators import support_torch_compile +from vllm.config import VllmConfig +from vllm.distributed import get_pp_group +from vllm.logger import init_logger +from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY +from vllm.model_executor.layers.linear import ( + ColumnParallelLinear, + ReplicatedLinear, + RowParallelLinear, +) +from vllm.model_executor.layers.logits_processor import LogitsProcessor +from vllm.model_executor.layers.quantization import QuantizationConfig +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead +from vllm.model_executor.model_loader.weight_utils import default_weight_loader +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.qwen2_5_omni_thinker import ( + Qwen2_5OmniConditionalGenerationMixin, + Qwen2_5OmniThinkerDummyInputsBuilder, + Qwen2_5OmniThinkerMultiModalProcessor, +) +from vllm.model_executor.models.qwen2_5_vl import ( + Qwen2_5_VisionAttention, + Qwen2_5_VisionRotaryEmbedding, + Qwen2_5_VLProcessingInfo, +) +from vllm.model_executor.models.qwen2_audio import Qwen2AudioInputs, Qwen2AudioProcessingInfo +from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM +from vllm.model_executor.models.utils import ( + AutoWeightsLoader, + WeightsMapper, + _merge_multimodal_embeddings, + flatten_bn, + maybe_prefix, +) +from vllm.model_executor.models.vision import ( + get_vit_attn_backend, +) +from vllm.multimodal import MULTIMODAL_REGISTRY +from vllm.multimodal.inputs import ( + MultiModalKwargsItems, +) +from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems +from vllm.multimodal.processing import ( + MultiModalPromptUpdates, + PlaceholderFeaturesInfo, + PromptReplacement, + PromptUpdate, + PromptUpdateDetails, +) +from vllm.platforms import _Backend +from vllm.sequence import IntermediateTensors +from vllm.v1.sample.metadata import SamplingMetadata + +from .qwen3_moe import Qwen3MoeModel + +try: + import flash_attn +except (ImportError, ModuleNotFoundError): + flash_attn = None + +from torch import Tensor + +logger = init_logger(__name__) + + +# Local helper functions (not available in vLLM) +def conv3d_to_linear_weight(weight: torch.Tensor) -> torch.Tensor: + """Convert 3D convolutional weight to linear weight for patch embedding.""" + # Weight shape: (out_channels, in_channels, t, h, w) + # Need to flatten spatial dimensions: (out_channels, in_channels * t * h * w) + flattened_weight = weight.flatten(1) # Flatten all dims except out_channels + return flattened_weight + + +def get_llm_pos_ids_for_vision( + start_idx: int, + vision_idx: int, + spatial_merge_size: int, + t_index: list[torch.Tensor], + grid_hs: list[torch.Tensor], + grid_ws: list[torch.Tensor], +) -> torch.Tensor: + """Generate 3D position IDs for vision tokens (temporal, height, width).""" + llm_grid_h = grid_hs[vision_idx] // spatial_merge_size + llm_grid_w = grid_ws[vision_idx] // spatial_merge_size + h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten().float() + w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten().float() + t_index_tensor = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().float() + llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index]) + llm_pos_ids = llm_pos_ids + start_idx + return llm_pos_ids + + +# Type aliases for audio inputs +Qwen2_5OmniAudioFeatureInputs = Qwen2AudioInputs + + +def _get_feat_extract_output_lengths(input_lengths: torch.Tensor): + input_lengths_leave = input_lengths % 100 + feat_lengths = (input_lengths_leave - 1) // 2 + 1 + output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 + return feat_lengths, output_lengths + + +class Qwen3_VisionPatchEmbed(nn.Module): + def __init__( + self, + patch_size: int = 14, + temporal_patch_size: int = 2, + in_channels: int = 3, + hidden_size: int = 1152, + ) -> None: + super().__init__() + self.patch_size = patch_size + self.temporal_patch_size = temporal_patch_size + self.hidden_size = hidden_size + + kernel_size = (temporal_patch_size, patch_size, patch_size) + self.proj = ReplicatedLinear( + in_channels * math.prod(kernel_size), + hidden_size, + bias=True, + return_bias=False, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + L, C = x.shape + x = self.proj(x) + return x + + +class Qwen3_VisionMLP(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: int, + bias: bool = False, + act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ): + super().__init__() + self.linear_fc1 = ColumnParallelLinear( + in_features, + hidden_features, + bias=bias, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.linear_fc1", + ) + self.linear_fc2 = RowParallelLinear( + hidden_features, + in_features, + bias=bias, + quant_config=quant_config, + return_bias=False, + prefix=f"{prefix}.linear_fc2", + ) + self.act_fn = act_fn + + def forward(self, x: torch.Tensor): + mlp_output = self.linear_fc2(self.act_fn(self.linear_fc1(x))) + return mlp_output + + +class Qwen3_VisionBlock(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_hidden_dim: int, + act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu, + norm_layer: Callable[[int], nn.Module] | None = None, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + self.norm1 = norm_layer(dim) + self.norm2 = norm_layer(dim) + self.attn = Qwen2_5_VisionAttention( + embed_dim=dim, + num_heads=num_heads, + projection_size=dim, + quant_config=quant_config, + prefix=f"{prefix}.attn", + ) + self.mlp = Qwen3_VisionMLP( + dim, + mlp_hidden_dim, + act_fn=act_fn, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.mlp", + ) + + def forward( + self, + x: torch.Tensor, + cu_seqlens: torch.Tensor, + rotary_pos_emb: torch.Tensor, + max_seqlen: int | None = None, # Only used for Flash Attention + seqlens: list[int] | None = None, # Only used for xFormers + ) -> torch.Tensor: + x = x + self.attn( + self.norm1(x), + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + max_seqlen=max_seqlen, + seqlens=seqlens, + ) + + x = x + self.mlp(self.norm2(x)) + return x + + +class Qwen3_VisionPatchMerger(nn.Module): + def __init__( + self, + d_model: int, + context_dim: int, + norm_layer: Callable[[int], nn.Module] | None = None, + spatial_merge_size: int = 2, + use_postshuffle_norm: bool = False, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + ) -> None: + super().__init__() + self.hidden_size = context_dim * (spatial_merge_size**2) + + self.use_postshuffle_norm = use_postshuffle_norm + if self.use_postshuffle_norm: + context_dim = self.hidden_size + + if norm_layer is None: + norm_layer = partial(nn.LayerNorm, eps=1e-6) + self.use_postshuffle_norm = use_postshuffle_norm + self.ln_q = norm_layer(self.hidden_size if use_postshuffle_norm else context_dim) + self.mlp = nn.ModuleList( + [ + ColumnParallelLinear( + self.hidden_size, + self.hidden_size, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.mlp.0", + ), + nn.GELU(), + RowParallelLinear( + self.hidden_size, + d_model, + bias=True, + quant_config=quant_config, + prefix=f"{prefix}.mlp.2", + ), + ] + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.use_postshuffle_norm: + x = self.ln_q(x.view(-1, self.hidden_size)) + else: + x = self.ln_q(x).view(-1, self.hidden_size) + + mlp_fc1, mlp_act, mlp_fc2 = self.mlp + x_parallel, _ = mlp_fc1(x) + x_parallel = mlp_act(x_parallel) + out, _ = mlp_fc2(x_parallel) + return out + + +class Qwen3Omni_VisionTransformer(nn.Module): + def __init__( + self, + vision_config, + norm_eps: float = 1e-6, + quant_config: QuantizationConfig | None = None, + prefix: str = "", + attn_backend_override: _Backend | None = None, + ) -> None: + super().__init__() + self.hidden_size = vision_config.hidden_size + self.num_heads = vision_config.num_heads + self.image_size = vision_config.image_size + self.patch_size = vision_config.patch_size + self.spatial_merge_size = vision_config.spatial_merge_size + self.spatial_merge_unit = self.spatial_merge_size**2 + self.temporal_patch_size = vision_config.temporal_patch_size + self.num_grid_per_side = self.image_size // self.patch_size + self.apply_vit_abs_pos_embed = vision_config.apply_vit_abs_pos_embed + self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes + + self.patch_embed = Qwen3_VisionPatchEmbed( + patch_size=self.patch_size, + temporal_patch_size=self.temporal_patch_size, + in_channels=vision_config.in_channels, + hidden_size=self.hidden_size, + ) + + # vit pos embedding, TODO: spatial_patch_size vs patch_size + if self.apply_vit_abs_pos_embed: + self.pos_embed = nn.Embedding(self.num_grid_per_side**2, self.hidden_size) + else: + self.pos_embed = nn.Parameter(torch.empty([1, self.num_grid_per_side**2, self.hidden_size])) + + norm_layer = partial(nn.LayerNorm, eps=norm_eps) + head_dim = self.hidden_size // self.num_heads + self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) + + self.blocks = nn.ModuleList( + [ + Qwen3_VisionBlock( + dim=self.hidden_size, + num_heads=self.num_heads, + mlp_hidden_dim=vision_config.intermediate_size, + act_fn=_ACTIVATION_REGISTRY[vision_config.hidden_act], + norm_layer=norm_layer, + quant_config=quant_config, + prefix=f"{prefix}.blocks.{layer_idx}", + ) + for layer_idx in range(vision_config.depth) + ] + ) + self.merger = Qwen3_VisionPatchMerger( + d_model=vision_config.out_hidden_size, + context_dim=self.hidden_size, + norm_layer=norm_layer, + spatial_merge_size=self.spatial_merge_size, + quant_config=quant_config, + prefix=f"{prefix}.merger", + ) + if self.deepstack_visual_indexes is not None: + self.merger_list = nn.ModuleList( + [ + Qwen3_VisionPatchMerger( + d_model=vision_config.out_hidden_size, + context_dim=self.hidden_size, + spatial_merge_size=self.spatial_merge_size, + use_postshuffle_norm=True, + norm_layer=norm_layer, + quant_config=quant_config, + prefix=f"{prefix}.merger_list.{layer_idx}", + ) + for layer_idx in range(len(self.deepstack_visual_indexes)) + ] + ) + + self.attn_backend = get_vit_attn_backend( + head_size=head_dim, + dtype=torch.get_default_dtype(), + # attn_backend_override=attn_backend_override, + ) + if self.attn_backend != _Backend.FLASH_ATTN and get_flash_attn_version() is not None: + self.attn_backend = _Backend.FLASH_ATTN + + @property + def dtype(self) -> torch.dtype: + return self.patch_embed.proj.weight.dtype + + @property + def device(self) -> torch.device: + return self.patch_embed.proj.weight.device + + def rot_pos_emb(self, grid_thw): + pos_ids = [] + for t, h, w in grid_thw: + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) + hpos_ids = hpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + hpos_ids = hpos_ids.permute(0, 2, 1, 3) + hpos_ids = hpos_ids.flatten() + + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) + wpos_ids = wpos_ids.reshape( + h // self.spatial_merge_size, + self.spatial_merge_size, + w // self.spatial_merge_size, + self.spatial_merge_size, + ) + wpos_ids = wpos_ids.permute(0, 2, 1, 3) + wpos_ids = wpos_ids.flatten() + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) + pos_ids = torch.cat(pos_ids, dim=0) + max_grid_size = grid_thw[:, 1:].max() + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) + return rotary_pos_emb + + def fast_pos_embed_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor: + num_grid_per_side = self.num_grid_per_side + m_size = self.spatial_merge_size + hidden_dim = self.pos_embed.embedding_dim + + outputs = [] + for t, h, w in grid_thw: + h_idxs = torch.linspace(0, num_grid_per_side - 1, h, dtype=torch.float32, device=self.device) + w_idxs = torch.linspace(0, num_grid_per_side - 1, w, dtype=torch.float32, device=self.device) + + h_floor = h_idxs.to(torch.long) + w_floor = w_idxs.to(torch.long) + h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1) + w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1) + + dh = h_idxs - h_floor + dw = w_idxs - w_floor + + # Create meshgrid view for all h, w vars + dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij") + h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij") + h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij") + h_floor_grid_idx = h_floor_grid * num_grid_per_side + h_ceil_grid_idx = h_ceil_grid * num_grid_per_side + + # original computation of weights + # w00 = (1 - dh_grid) * (1 - dw_grid) + # w01 = (1 - dh_grid) * dw_grid + # w10 = dh_grid * (1 - dw_grid) + # w11 = dh_grid * dw_grid + # we reuse w11 here to avoid duplicate + # dh_grid * dw_grid computation + w11 = dh_grid * dw_grid + w10 = dh_grid - w11 + w01 = dw_grid - w11 + w00 = 1 - dh_grid - dw_grid + w11 + + idx00 = h_floor_grid_idx + w_floor_grid + idx01 = h_floor_grid_idx + w_ceil_grid + idx10 = h_ceil_grid_idx + w_floor_grid + idx11 = h_ceil_grid_idx + w_ceil_grid + + indices = torch.stack([idx00, idx01, idx10, idx11], dim=0).reshape(4, -1) + weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1) + weights = weights.to(dtype=self.dtype, device=self.device) + + embeds = self.pos_embed(indices) + weighted_embeds = embeds * weights + p0, p1, p2, p3 = weighted_embeds.unbind(dim=0) + combined = p0 + p1 + p2 + p3 + + combined = combined.view(h * w, hidden_dim) + repeated = combined.unsqueeze(0).expand(t, -1, -1).contiguous() + repeated = repeated.view(t, h // m_size, m_size, w // m_size, m_size, hidden_dim) + repeated = repeated.permute(0, 1, 3, 2, 4, 5).reshape(-1, hidden_dim) + outputs.append(repeated) + + return torch.cat(outputs, dim=0) + + def compute_attn_mask_seqlen( + self, + cu_seqlens: torch.Tensor, + ) -> tuple[int | None, list[int] | None]: + max_seqlen, seqlens = None, None + if self.attn_backend == _Backend.FLASH_ATTN: + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() + elif self.attn_backend == _Backend.XFORMERS: + seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() + return max_seqlen, seqlens + + def forward( + self, + x: torch.Tensor, + grid_thw: list[list[int]], + ) -> torch.Tensor: + hidden_states = x.to(device=self.device, dtype=self.dtype) + hidden_states = self.patch_embed(hidden_states) + + if self.apply_vit_abs_pos_embed: + pos_embeds = self.fast_pos_embed_interpolate(grid_thw) + hidden_states = hidden_states + pos_embeds + rotary_pos_emb = self.rot_pos_emb(grid_thw) + + cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( + dim=0, + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, + ) + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) + + hidden_states = hidden_states.unsqueeze(1) + rotary_pos_emb = rotary_pos_emb.to(hidden_states.device) + max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens) + + hidden_states_list = [] + deepstack_visual_indexes = self.deepstack_visual_indexes + + for layer_num, blk in enumerate(self.blocks): + hidden_states = blk( + hidden_states, + cu_seqlens=cu_seqlens, + rotary_pos_emb=rotary_pos_emb, + max_seqlen=max_seqlen, + seqlens=seqlens, + ) + if deepstack_visual_indexes is not None and layer_num in deepstack_visual_indexes: + hidden_states_list.append(hidden_states) + + hidden_states = self.merger(hidden_states) + + # processing deepstack + if deepstack_visual_indexes is not None: + processed_hidden_states_list = [hidden_states] + for idx, x in enumerate(hidden_states_list): + x = self.merger_list[idx](x) + processed_hidden_states_list.append(x) + # we cat the original visual features and deepstack features + # along the feature dim + hidden_states = torch.cat( + processed_hidden_states_list, dim=1 + ) # [seq_len, hidden_size * (1 + depth_of_deepstack)] + + return hidden_states + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + stacked_params_mapping = [ + # (param_name, shard_name, shard_id) + ("attn.qkv.", "attn.q.", "q"), + ("attn.qkv.", "attn.k.", "k"), + ("attn.qkv.", "attn.v.", "v"), + ] + params_dict = dict(self.named_parameters(remove_duplicate=False)) + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + if name.endswith("patch_embed.proj.weight"): + loaded_weight = conv3d_to_linear_weight(loaded_weight) + + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + name = name.replace(weight_name, param_name) + + param = params_dict[name] + weight_loader = param.weight_loader + weight_loader(param, loaded_weight, shard_id) + break + else: + 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 + + +@support_torch_compile( + dynamic_arg_dims={ + "input_ids": 0, + "positions": -1, + "intermediate_tensors": 0, + "inputs_embeds": 0, + "deepstack_input_embeds": 0, + } +) +class Qwen3MoeLLMModel(Qwen3MoeModel): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__(vllm_config=vllm_config, prefix=prefix) + # from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeTextRMSNorm + # config = vllm_config.model_config.hf_config + # self.norm = Qwen3OmniMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.deepstack_multiscale_layer_start = 1 + + def _deepstack_process(self, hidden_states: torch.Tensor, visual_embeds: torch.Tensor) -> torch.Tensor: + """ + Process deepstack visual embeddings by selectively adding them only to visual positions. + This matches the transformers implementation which uses visual_pos_masks. + + Args: + hidden_states: Hidden states tensor [num_tokens, hidden_size] + visual_embeds: Visual embeddings tensor [num_tokens, hidden_size] (zero-padded for non-visual positions) + + Returns: + Modified hidden_states with visual embeddings added only to visual positions + """ + # Derive visual position mask from visual_embeds (non-zero positions are visual) + # Use a small threshold to handle floating point precision issues + visual_pos_masks = torch.any(torch.abs(visual_embeds) > 1e-8, dim=-1) + + if visual_pos_masks.any(): + visual_pos_masks = visual_pos_masks.to(hidden_states.device) + visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype) + # Match transformers implementation EXACTLY (lines 2966-2969): + # In transformers: local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds + # Note: visual_embeds in transformers is NOT indexed, meaning it's already shaped + # [num_visual_tokens, hidden_size]. In vllm-omni, visual_embeds is [num_tokens, hidden_size] + # zero-padded, so we extract only visual positions to match transformers behavior. + visual_embeds_only = visual_embeds[visual_pos_masks, :] # Extract only visual positions + local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds_only + hidden_states[visual_pos_masks, :] = local_this + + return hidden_states + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + *, + capture_layer_indices: Sequence[int] | None = None, + return_hidden_states: bool = False, + deepstack_input_embeds: IntermediateTensors | None = None, + ) -> torch.Tensor | IntermediateTensors: + if get_pp_group().is_first_rank: + if inputs_embeds is not None: + hidden_states = inputs_embeds + else: + hidden_states = self.get_input_embeddings(input_ids) + residual = None + else: + assert intermediate_tensors is not None + hidden_states = intermediate_tensors["hidden_states"] + residual = intermediate_tensors["residual"] + capture_set = set(capture_layer_indices) if capture_layer_indices else None + captured_hidden_states: dict[str, torch.Tensor] | None = {} if return_hidden_states else None + + for layer_idx, layer in enumerate(self.layers[self.start_layer : self.end_layer]): + layer_idx = layer_idx + self.start_layer + + if captured_hidden_states is not None and capture_set is not None: + if layer_idx in capture_set: + captured_hidden_states[str(layer_idx)] = hidden_states.clone().view(-1, hidden_states.shape[-1]) + + hidden_states, residual = layer( + positions, + hidden_states, + residual, + ) + + if deepstack_input_embeds is not None and layer_idx in range(0, len(deepstack_input_embeds)): + # Use _deepstack_process to match transformers implementation + # This selectively adds visual embeddings only to visual positions + hidden_states = self._deepstack_process( + hidden_states, deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"] + ) + + if not get_pp_group().is_last_rank: + return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) + hidden_states, _ = self.norm(hidden_states, residual) + if captured_hidden_states is not None: + return hidden_states, captured_hidden_states + else: + return hidden_states, None + + +class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM): + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super(Qwen3MoeForCausalLM, self).__init__() + config = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + self.config = config + self.quant_config = quant_config + self.model = Qwen3MoeLLMModel(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) + self.lm_head = ParallelLMHead(config.vocab_size, config.hidden_size, quant_config=quant_config) + if self.config.tie_word_embeddings: + self.lm_head.weight = self.model.embed_tokens.weight + self.logits_processor = LogitsProcessor(config.vocab_size) + self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors + + +class Qwen3OmniMoeThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo): + def get_hf_config(self): + return self.ctx.get_hf_config(Qwen3OmniMoeConfig).thinker_config + + def get_hf_processor(self, **kwargs: object) -> Qwen3OmniMoeProcessor: + processor = self.ctx.get_hf_processor( + Qwen3OmniMoeProcessor, + use_fast=kwargs.pop("use_fast", True), + **kwargs, + ) + if not hasattr(processor, "audio_token"): + processor.audio_token = "<|audio_pad|>" + if not hasattr(processor, "image_token"): + processor.image_token = "<|image_pad|>" + if not hasattr(processor, "video_token"): + processor.video_token = "<|video_pad|>" + return processor + + def get_feature_extractor(self, **kwargs: object): + hf_processor = self.get_hf_processor(**kwargs) + feature_extractor = hf_processor.feature_extractor # type: ignore + assert isinstance(feature_extractor, WhisperFeatureExtractor) + return feature_extractor + + def get_supported_mm_limits(self) -> Mapping[str, int | None]: + return {"audio": None, "image": None, "video": None} + + +Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder + + +class Qwen3OmniMoeThinkerMultiModalProcessor( + Qwen2_5OmniThinkerMultiModalProcessor, +): + def _call_hf_processor( + self, + prompt: str, + mm_data: Mapping[str, object], + mm_kwargs: Mapping[str, object], + tok_kwargs: Mapping[str, object], + ) -> BatchFeature: + mm_data = dict(mm_data) + audios = mm_data.pop("audios", []) + + def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray: + length = x.shape[-1] + if length % hop_length != 0: + pad_length = hop_length - (length % hop_length) + x = np.pad(x, (0, pad_length), mode="constant", constant_values=0) + return x + + # NOTE: WhisperFeatureExtractor cannot handle empty list of audios + feature_extractor = self.info.get_feature_extractor() + hop_length = feature_extractor.hop_length + if audios: + # NOTE: Qwen3-Omni processor accept "audio" + # To make sure the cache works with padding=True, we pre-padded + # the audio to multiple of hop_length. + mm_data["audio"] = [ + pad_to_hop_length(audio, hop_length) + if isinstance(audio, np.ndarray) + else (pad_to_hop_length(audio[0], hop_length), audio[1]) + for audio in audios + ] + + # TODO(Isotr0py): Remove this patch after upstream fix PR + # released and Transformers version update: + # https://github.com/huggingface/transformers/pull/41473 + mm_kwargs = dict(mm_kwargs) + tok_kwargs = dict(tok_kwargs) + if Version(TRANSFORMERS_VERSION) < Version("4.58.0"): + # move truncation to audio_kwargs level to avoid conflict + # with tok_kwargs + mm_kwargs["audio_kwargs"] = {"truncation": mm_kwargs.pop("truncation", False)} + mm_kwargs["text_kwargs"] = {"truncation": tok_kwargs.pop("truncation", False)} + + hf_inputs = super()._call_hf_processor( + prompt=prompt, + mm_data=mm_data, + mm_kwargs=mm_kwargs, + tok_kwargs=tok_kwargs, + ) + + if ( + "audio_feature_lengths" in hf_inputs + and "feature_attention_mask" in hf_inputs + and (audios := mm_data.get("audio", [])) + ): + audio_num_frames = [] + for _, audio in enumerate(audios): + audio_length = len(audio[0]) if isinstance(audio, tuple) else len(audio) + num_frame = ( + (audio_length // hop_length) if audio_length % hop_length == 0 else (audio_length // hop_length - 1) + ) + if mm_kwargs.get("truncation", False): + num_frame = min(num_frame, feature_extractor.n_samples // hop_length) + audio_num_frames.append(num_frame) + hf_inputs["feature_attention_mask"] = [torch.ones(num_frame) for num_frame in audio_num_frames] + hf_inputs["audio_feature_lengths"] = torch.tensor(audio_num_frames) + return hf_inputs + + def _maybe_apply_prompt_updates( + self, + mm_items: MultiModalDataItems, + prompt_ids: list[int], + mm_kwargs: MultiModalKwargsItems, + mm_prompt_updates: MultiModalPromptUpdates, + is_update_applied: bool, + ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]: + """ + Qwen3-Omni reimplements this function to handle `use_audio_in_video`. + """ + mm_item_counts = mm_items.get_all_counts() + self._validate_mm_kwargs(mm_kwargs, mm_item_counts) + + use_audio_in_video = False + if "video" in mm_kwargs: + for item in mm_kwargs["video"]: + if item and item["use_audio_in_video"].data: + use_audio_in_video = True + else: + use_audio_in_video = False + + # normal case with `use_audio_in_video=False` + if is_update_applied: + mm_placeholders = self._find_mm_placeholders( + prompt_ids, + mm_prompt_updates, + ) + self._validate_mm_placeholders( + mm_placeholders, + mm_item_counts, + ) + else: + if use_audio_in_video and "audio" in mm_prompt_updates: + filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"} + prompt_ids, mm_placeholders = self._apply_prompt_updates( + prompt_ids, + filtered_updates, + ) + # Derive audio placeholders from video placeholders + mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates) + else: + prompt_ids, prompt, mm_placeholders = self._apply_prompt_updates( + prompt_ids, + mm_prompt_updates, + ) + + self._validate_mm_placeholders( + mm_placeholders, + mm_item_counts, + ) + + return prompt_ids, prompt, mm_placeholders + + def get_updates_use_audio_in_video( + self, + thinker_config: PretrainedConfig, + audio_len: int, + video_grid_thw: list[int] | torch.Tensor, + video_second_per_grid_t: float, + ) -> list[int]: + shift = 0 + audio_token_id = thinker_config.audio_token_id + video_token_id = thinker_config.video_token_id + audio_start_token_id = thinker_config.audio_start_token_id + audio_end_token_id = thinker_config.audio_end_token_id + spatial_merge_size = thinker_config.vision_config.spatial_merge_size + position_id_per_seconds = thinker_config.position_id_per_seconds + audio_token_indices = np.arange(next(iter([audio_len]))) + curr_video_grid_thw = next(iter([video_grid_thw])) + height = curr_video_grid_thw[1] // spatial_merge_size + width = curr_video_grid_thw[2] // spatial_merge_size + video_token_indices = np.arange(curr_video_grid_thw[0]).reshape(-1, 1, 1) + video_token_indices = np.broadcast_to( + video_token_indices, (video_token_indices.shape[0], height, width) + ).reshape(-1) + video_token_indices = ( + (video_token_indices + shift) * next(iter([video_second_per_grid_t])) * position_id_per_seconds + ) + video_data_index, audio_data_index = 0, 0 + updates = [audio_start_token_id] + while video_data_index < len(video_token_indices) and audio_data_index < len(audio_token_indices): + if video_token_indices[video_data_index] <= audio_token_indices[audio_data_index]: + updates += [video_token_id] + video_data_index += 1 + else: + updates += [audio_token_id] + audio_data_index += 1 + if video_data_index < len(video_token_indices): + updates += [video_token_id] * (len(video_token_indices) - video_data_index) + if audio_data_index < len(audio_token_indices): + updates += [audio_token_id] * (len(audio_token_indices) - audio_data_index) + updates += [audio_end_token_id] + return updates + + def _get_prompt_updates( + self, + mm_items: MultiModalDataItems, + hf_processor_mm_kwargs: Mapping[str, Any], + out_mm_kwargs: MultiModalKwargsItems, + ) -> Sequence[PromptUpdate]: + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) + tokenizer = self.info.get_tokenizer() + image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs) + vocab = tokenizer.get_vocab() + + audio_token = processor.audio_token + image_token = processor.image_token + video_token = processor.video_token + audio_token_id = vocab[audio_token] + image_token_id = vocab[image_token] + video_token_id = vocab[video_token] + + out_mm_data = out_mm_kwargs.get_data() + audio_feature_lengths = out_mm_data.get("audio_feature_lengths") + feature_attention_mask = out_mm_data.get("feature_attention_mask") + if audio_feature_lengths is None and feature_attention_mask is None: + audio_output_lengths = [] + elif audio_feature_lengths is not None: + _, audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths) + audio_output_lengths = audio_output_lens.tolist() + elif feature_attention_mask is not None: + assert isinstance(feature_attention_mask, torch.Tensor) + _, audio_output_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1)) + audio_output_lengths = audio_output_lens.tolist() + + # number of audios read from video. + audio_in_video_item_idx = 0 + audio_item_idx = 0 + + def get_replacement_qwen2_audio(item_idx: int): + nonlocal audio_item_idx + item_idx += audio_in_video_item_idx + + audio_item_idx += 1 + + num_features = audio_output_lengths[item_idx] + if num_features == 0: + audios = mm_items.get_items("audio", AudioProcessorItems) + audio = audios.get(item_idx) + raise ValueError( + f"The audio {audio} (len={len(audio)}) is too short to be represented inside the model" + ) + + return [audio_token_id] * num_features + + def get_replacement_qwen2_vision(item_idx: int, modality: str): + grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx] + # grid_thw = out_mm_kwargs[f"{modality}_grid_thw"][item_idx] + assert isinstance(grid_thw, torch.Tensor) + merge_length = image_processor.merge_size**2 + + token_id = image_token_id if modality == "image" else video_token_id + return [token_id] * (int(grid_thw.prod()) // merge_length) + + use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False) + thinker_config = self.info.get_hf_config() + + def get_replacement_qwen2_use_audio_in_video(item_idx: int): + nonlocal audio_in_video_item_idx + # audio_num_features = audio_output_lengths[audio_item_idx + item_idx] + # video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx] + audio_num_features = audio_output_lengths[audio_in_video_item_idx + item_idx] + video_grid_thw = out_mm_data["video_grid_thw"][item_idx] + + audio_in_video_item_idx += 1 + + second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None) + if second_per_grid_ts: + video_second_per_grid_t = second_per_grid_ts[item_idx] + else: + video_second_per_grid_t = 2.0 + + placeholder = self.get_updates_use_audio_in_video( + thinker_config=thinker_config, + audio_len=audio_num_features, + video_grid_thw=video_grid_thw, + video_second_per_grid_t=video_second_per_grid_t, + ) + return PromptUpdateDetails.select_token_id(placeholder, embed_token_id=video_token_id) + + video_replacement_fn = ( + get_replacement_qwen2_use_audio_in_video + if use_audio_in_video + else partial(get_replacement_qwen2_vision, modality="video") + ) + + return [ + PromptReplacement( + modality="audio", + target=audio_token, + replacement=get_replacement_qwen2_audio, + ), + PromptReplacement( + modality="image", + target=image_token, + replacement=partial(get_replacement_qwen2_vision, modality="image"), + ), + PromptReplacement( + modality="video", + target=video_token, + replacement=video_replacement_fn, + ), + ] + + def _derive_audio_from_video_placeholders( + self, + placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], + mm_prompt_updates: Mapping[str, list[PromptUpdateDetails]], + ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: + """ + Helper to derive audio placeholders from video placeholders when + use_audio_in_video=True. + """ + if "video" not in placeholders: + return placeholders + + # Validate audio and video counts match + num_videos = len(placeholders["video"]) + num_audios = len(mm_prompt_updates.get("audio", [])) + if num_audios != num_videos: + raise ValueError( + f"use_audio_in_video requires equal number of audio and video items, got {num_audios=}, {num_videos=}" + ) + + tokenizer = self.info.get_tokenizer() + processor = self.info.get_hf_processor() + audio_token_id = tokenizer.get_vocab()[processor.audio_token] + + result_placeholders = dict(placeholders) + audio_placeholders = [] + + # Each video is paired with one audio + for video_idx, video_placeholder in enumerate(placeholders["video"]): + # Create is_embed mask selecting only audio tokens + audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id + + audio_placeholder = PlaceholderFeaturesInfo( + modality="audio", + item_idx=video_idx, + start_idx=video_placeholder.start_idx, + tokens=video_placeholder.tokens, + is_embed=audio_is_embed, + ) + audio_placeholders.append(audio_placeholder) + + result_placeholders["audio"] = audio_placeholders + return result_placeholders + + def _get_raw_input_ids( + self, + token_ids: list[int], + use_audio_in_video: bool = False, + ) -> list[int]: + tokenizer = self.info.get_tokenizer() + vision_bos_token = tokenizer.encode(tokenizer.vision_bos_token)[0] + vision_eos_token = tokenizer.encode(tokenizer.vision_eos_token)[0] + audio_bos_token = tokenizer.encode(tokenizer.audio_bos_token)[0] + audio_eos_token = tokenizer.encode(tokenizer.audio_eos_token)[0] + audio_token = tokenizer.encode("<|audio_pad|>")[0] + image_token = tokenizer.encode("<|image_pad|>")[0] + video_token = tokenizer.encode("<|video_pad|>")[0] + + result = token_ids[:] + if use_audio_in_video: + while True: + start = None + for i in range(len(result) - 1): + if result[i : i + 2] == [vision_bos_token, audio_bos_token]: + start = i + break + if start is not None: + end = None + for i in range(start + 2, len(result) - 1): + if result[i : i + 2] == [audio_eos_token, vision_eos_token]: + end = i + break + if end is not None: + result = result[:start] + [vision_bos_token, video_token, vision_eos_token] + result[end + 2 :] + else: + break + + for mm_token in [audio_token, image_token, video_token]: + compressed = [] + for x in result: + if x != mm_token or (not compressed or compressed[-1] != mm_token): + compressed.append(x) + result = compressed + + return result + + +class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixin): + def _process_audio_input( + self, + audio_input: Qwen2_5OmniAudioFeatureInputs, + audio_hashes: list[str] | None = None, + cached_audio_features: torch.Tensor | None = None, + ) -> torch.Tensor: + input_features = audio_input["input_features"] + audio_feature_lengths = audio_input["audio_feature_lengths"] + + audio_feature_lengths = flatten_bn(audio_feature_lengths, concat=True) + + audio_feat_lengths, audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths) + + audio_outputs = self.audio_tower( + input_features.to(self.audio_tower.dtype), + feature_lens=audio_feature_lengths, + aftercnn_lens=audio_feat_lengths, + ) + audio_features = audio_outputs.last_hidden_state + return audio_features.split(audio_output_lengths.tolist()) + + +@MULTIMODAL_REGISTRY.register_processor( + Qwen3OmniMoeThinkerMultiModalProcessor, + info=Qwen3OmniMoeThinkerProcessingInfo, + dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder, +) +class Qwen3OmniMoeThinkerForConditionalGeneration( + nn.Module, + SupportsMultiModal, + SupportsPP, + Qwen3OmniMoeConditionalGenerationMixin, +): + merge_by_field_config = True + + hf_to_vllm_mapper = WeightsMapper( + orig_to_new_prefix={ + "thinker.lm_head.": "language_model.lm_head.", + "thinker.model.": "language_model.model.", + "thinker.": "", + } + ) + + @classmethod + def get_placeholder_str(cls, modality: str, i: int) -> str | None: + if modality.startswith("image"): + return "<|vision_start|><|image_pad|><|vision_end|>" + if modality.startswith("video"): + return "<|vision_start|><|video_pad|><|vision_end|>" + if modality.startswith("audio"): + return "<|audio_start|><|audio_pad|><|audio_end|>" + + raise ValueError("Only image, video or audio modality is supported") + + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): + super().__init__() + thinker_config: Qwen3OmniMoeThinkerConfig = vllm_config.model_config.hf_config + quant_config = vllm_config.quant_config + multimodal_config = vllm_config.model_config.multimodal_config + self.config = thinker_config + self.multimodal_config = multimodal_config + + # force "use_flash_attention_2=True" to audio tower to align + # the results. + if flash_attn is not None: + audio_config = thinker_config.audio_config + audio_config._attn_implementation_autoset = True + audio_config._attn_implementation = "flash_attention_2" + else: + logger.warning( + "flash_attn is not available, the model may not yield the " + "exactly same result as the transformers implementation " + "in the audio tower part." + ) + + self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config) + + # attn_backend_override = ( + # multimodal_config.mm_encoder_attn_backend + # if multimodal_config is not None + # else None + # ) + self.visual = Qwen3Omni_VisionTransformer( + vision_config=thinker_config.vision_config, + norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), + quant_config=quant_config, + prefix=maybe_prefix(prefix, "visual"), + # attn_backend_override=attn_backend_override, + ) + self.quant_config = quant_config + + lm_vllm_config = vllm_config.with_hf_config(thinker_config.text_config, architectures=["Qwen3MoeForCausalLM"]) + + self.language_model = Qwen3MoeLLMForCausalLM( + vllm_config=lm_vllm_config, + prefix=maybe_prefix(prefix, "language_model"), + ) + + self.make_empty_intermediate_tensors = self.language_model.make_empty_intermediate_tensors + + self.use_deepstack = hasattr(thinker_config.vision_config, "deepstack_visual_indexes") + self.deepstack_num_level = ( + len(thinker_config.vision_config.deepstack_visual_indexes) if self.use_deepstack else 0 + ) + # register buffer for deepstack + self.deepstack_input_embeds = ( + [ + torch.zeros( + vllm_config.scheduler_config.max_num_batched_tokens, + thinker_config.text_config.hidden_size, + ) + for _ in range(self.deepstack_num_level) + ] + if self.use_deepstack + else None + ) + self.visual_dim = thinker_config.vision_config.out_hidden_size + self.multiscale_dim = self.visual_dim * self.deepstack_num_level + + def _get_deepstack_input_embeds(self, num_tokens: int) -> IntermediateTensors: + # get deepstack_input_embeds from buffer, and clear the buffer + return IntermediateTensors( + { + f"deepstack_input_embeds_{idx}": self.deepstack_input_embeds[idx][:num_tokens] + for idx in range(self.deepstack_num_level) + } + ) + + def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> None: + # set deepstack_input_embeds to buffer + num_tokens = deepstack_input_embeds.size(1) + if num_tokens > self.deepstack_input_embeds[0].size(0): + self.deepstack_input_embeds = [ + torch.zeros( + num_tokens, + self.config.text_config.hidden_size, + device=self.deepstack_input_embeds[0].device, + dtype=self.deepstack_input_embeds[0].dtype, + ) + for _ in range(self.deepstack_num_level) + ] + for idx in range(self.deepstack_num_level): + self.deepstack_input_embeds[idx][:num_tokens].copy_(deepstack_input_embeds[idx]) + + def _clear_deepstack_input_embeds(self, num_tokens: int) -> None: + # clear deepstack_input_embeds in buffer + if num_tokens > 0: + for idx in range(self.deepstack_num_level): + self.deepstack_input_embeds[idx][:num_tokens].zero_() + + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict: + mm_input_by_modality = {} + + # Preserve the order of modalities if there are multiple of them + # from the order of kwargs. + for input_key in kwargs: + if input_key in ("pixel_values", "image_embeds") and "image" not in mm_input_by_modality: + mm_input_by_modality["image"] = self._parse_and_validate_image_input(**kwargs) + if input_key in ("pixel_values_videos", "video_embeds") and "video" not in mm_input_by_modality: + mm_input_by_modality["video"] = self._parse_and_validate_video_input(**kwargs) + if input_key in ("input_audio_features") and "audio" not in mm_input_by_modality: + mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs) + return mm_input_by_modality + + def get_language_model(self) -> torch.nn.Module: + return self.language_model + + def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings | None: + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) + if not mm_input_by_modality: + return [] + + # The result multimodal_embeddings is tuple of tensors, with each + # tensor correspoending to a multimodal data item (image or video). + multimodal_embeddings: tuple[torch.Tensor, ...] = () + + # NOTE: It is important to iterate over the keys in this dictionary + # to preserve the order of the modalities. + for modality in mm_input_by_modality: + multimodal_input = mm_input_by_modality[modality] + if modality == "image": + image_embeddings = self._process_image_input(multimodal_input) + multimodal_embeddings += tuple(image_embeddings) + if modality == "video": + video_embeddings = self._process_video_input(multimodal_input) + multimodal_embeddings += tuple(video_embeddings) + if modality == "audio": + audio_embeddings = self._process_audio_input(multimodal_input) + multimodal_embeddings += tuple(audio_embeddings) + return multimodal_embeddings + + def _get_text_embeddings( + self, + input_ids: Tensor, + embed_input_ids: Callable[[Tensor], Tensor], + *, + is_multimodal: Tensor | None, + handle_oov_mm_token: bool, + ) -> Tensor: + if handle_oov_mm_token and is_multimodal is not None: + is_text = ~is_multimodal + text_embeds = embed_input_ids(input_ids[is_text]) + + return torch.empty( + (input_ids.shape[0], text_embeds.shape[1]), + dtype=text_embeds.dtype, + device=text_embeds.device, + ).masked_scatter_(is_text.unsqueeze_(-1), text_embeds) + + return embed_input_ids(input_ids) + + def get_input_embeddings( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + handle_oov_mm_token: bool = False, + ) -> torch.Tensor: + # Compute is_multimodal if not provided and we have multimodal embeddings + if is_multimodal is None and multimodal_embeddings is not None and len(multimodal_embeddings) > 0: + # Create a boolean tensor marking which positions are multimodal tokens + is_multimodal = ( + (input_ids == self.config.image_token_id) + | (input_ids == self.config.video_token_id) + | (input_ids == self.config.audio_token_id) + ) + + inputs_embeds = self._get_text_embeddings( + input_ids, + self.language_model.get_input_embeddings, + is_multimodal=is_multimodal, + handle_oov_mm_token=handle_oov_mm_token, + ) + + if multimodal_embeddings is None or len(multimodal_embeddings) == 0: + return inputs_embeds + + deepstack_input_embeds = None + # TODO (ywang96): support overlapping modalitiy embeddings so that + # `use_audio_in_video` will work on V1. + # split the feat dim to obtain multi-scale visual feature + has_vision_embeddings = [ + embeddings.shape[-1] != self.config.text_config.hidden_size for embeddings in multimodal_embeddings + ] + if self.visual.deepstack_visual_indexes is not None and any(has_vision_embeddings): + multiscale_len = len(self.visual.deepstack_visual_indexes) + multimodal_embeddings_multiscale = [] + is_vision = torch.zeros_like(is_multimodal) + mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0] + mm_position_idx = 0 + for index, embeddings in enumerate(multimodal_embeddings): + num_tokens = embeddings.shape[0] + current_positions = mm_positions[mm_position_idx : mm_position_idx + num_tokens] + + # Vision embeddings + if embeddings.shape[-1] != self.config.text_config.hidden_size: + visual_dim = embeddings.shape[-1] // (multiscale_len + 1) + multi_dim = visual_dim * multiscale_len + embeddings_main, embeddings_multiscale = torch.split(embeddings, [visual_dim, multi_dim], dim=-1) + multimodal_embeddings[index] = embeddings_main + multimodal_embeddings_multiscale.append(embeddings_multiscale) + is_vision[current_positions] = True + + # Audio embeddings + else: + is_vision[current_positions] = False + + mm_position_idx += num_tokens + + deepstack_input_embeds = inputs_embeds.new_zeros( + inputs_embeds.size(0), multiscale_len * inputs_embeds.size(1) + ) + deepstack_input_embeds = _merge_multimodal_embeddings( + inputs_embeds=deepstack_input_embeds, + multimodal_embeddings=multimodal_embeddings_multiscale, + is_multimodal=is_vision, + ) + deepstack_input_embeds = ( + deepstack_input_embeds.view(inputs_embeds.shape[0], multiscale_len, visual_dim) + .permute(1, 0, 2) + .contiguous() + ) + self._set_deepstack_input_embeds(deepstack_input_embeds) + + inputs_embeds = _merge_multimodal_embeddings( + inputs_embeds=inputs_embeds, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + ) + + return inputs_embeds + + def forward( + self, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + *, + capture_layer_indices: Sequence[int] | None = None, + return_hidden_states: bool = False, + **kwargs: object, + ) -> torch.Tensor | IntermediateTensors: + if intermediate_tensors is not None: + inputs_embeds = None + + if self.use_deepstack and inputs_embeds is not None and get_pp_group().is_first_rank: + deepstack_input_embeds = self._get_deepstack_input_embeds(inputs_embeds.size(0)) + else: + deepstack_input_embeds = None + + hidden_states, captured_hidden_states = self.language_model.model( + input_ids, + positions, + intermediate_tensors, + inputs_embeds=inputs_embeds, + capture_layer_indices=capture_layer_indices, + return_hidden_states=return_hidden_states, + # args for deepstack + deepstack_input_embeds=deepstack_input_embeds, + ) + + if inputs_embeds is not None and get_pp_group().is_first_rank: + self._clear_deepstack_input_embeds(inputs_embeds.size(0)) + + return hidden_states.unsqueeze(0), captured_hidden_states + + def compute_logits( + self, + hidden_states: torch.Tensor, + sampling_metadata: SamplingMetadata = None, + ) -> torch.Tensor | None: + return self.language_model.compute_logits(hidden_states) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + loader = AutoWeightsLoader( + self, + skip_prefixes=["talker.", "code2wav."], + ) + loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) + + # Log load summary + try: + total_bytes = 0 + for name, param in self.named_parameters(): + if param is not None and param.data is not None: + total_bytes += param.data.numel() * param.data.element_size() + device = next(self.parameters()).device + logger.info( + "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s", + self.__class__.__name__, + True, + total_bytes / (1024**2), + str(device), + ) + except Exception: + pass + + return loaded_weights + + def get_mrope_input_positions( + self, + input_tokens: list[int], + hf_config: PretrainedConfig, + image_grid_thw: list[list[int]] | torch.Tensor | None, + video_grid_thw: list[list[int]] | torch.Tensor | None, + second_per_grid_ts: list[float] | None = None, + context_len: int = 0, + seq_len: int | None = None, + audio_feature_lengths: torch.Tensor | None = None, + use_audio_in_video: bool = False, + ) -> tuple[torch.Tensor, int]: + config = hf_config.thinker_config + if isinstance(image_grid_thw, list): + image_grid_thw = torch.tensor(image_grid_thw) + if isinstance(video_grid_thw, list): + video_grid_thw = torch.tensor(video_grid_thw) + input_ids = torch.tensor(input_tokens) + if input_ids is None or input_ids.ndim != 1: + raise ValueError("_omni3_get_input_positions_tensor expects 1D input_ids") + + seq_len = input_ids.shape[0] + if audio_feature_lengths is not None and not isinstance(audio_feature_lengths, torch.Tensor): + audio_feature_lengths = torch.as_tensor(audio_feature_lengths, dtype=torch.long) + if not len(second_per_grid_ts) and len(video_grid_thw): + second_per_grid_ts = 2.0 + second_per_grids = torch.ones(len(video_grid_thw), dtype=torch.float32) * second_per_grid_ts + else: + second_per_grids = torch.tensor(second_per_grid_ts, dtype=torch.float32) + + spatial_merge_size = config.vision_config.spatial_merge_size + image_token_id = config.image_token_id + video_token_id = config.video_token_id + audio_token_id = config.audio_token_id + vision_start_token_id = config.vision_start_token_id + audio_start_token_id = config.audio_start_token_id + position_id_per_seconds = config.position_id_per_seconds + + vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1) + if vision_start_indices.numel() > 0: + vision_tokens = input_ids[vision_start_indices + 1] + else: + vision_tokens = input_ids.new_empty((0,), dtype=input_ids.dtype) + audio_nums = torch.sum(input_ids == audio_start_token_id) + image_nums = (vision_tokens == image_token_id).sum() + video_nums = ( + (vision_tokens == audio_start_token_id).sum() + if use_audio_in_video + else (vision_tokens == video_token_id).sum() + ) + + llm_pos_ids_list: list[torch.Tensor] = [] + st = 0 + image_idx = 0 + video_idx = 0 + audio_idx = 0 + remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums # noqa: E501 + multimodal_nums = image_nums + audio_nums if use_audio_in_video else image_nums + video_nums + audio_nums # noqa: E501 + + for _ in range(multimodal_nums): + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + if (image_token_id in input_tokens or video_token_id in input_tokens) and ( + remain_videos > 0 or remain_images > 0 + ): + ed_vision_start = input_tokens.index(vision_start_token_id, st) + else: + ed_vision_start = len(input_tokens) + 1 + if audio_token_id in input_tokens and remain_audios > 0: + ed_audio_start = input_tokens.index(audio_start_token_id, st) + else: + ed_audio_start = len(input_tokens) + 1 + min_ed = min(ed_vision_start, ed_audio_start) + + if min_ed == ed_audio_start: + text_len = min_ed - st + if text_len != 0: + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + bos_len = 1 + llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + _, audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx]) + llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx + llm_pos_ids_list.append(llm_pos_ids) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + eos_len = 1 + llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + st += text_len + bos_len + audio_len + eos_len + audio_idx += 1 + remain_audios -= 1 + elif min_ed == ed_vision_start and input_ids[ed_vision_start + 1] == image_token_id: + text_len = min_ed - st + if text_len != 0: + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + bos_len = 1 + llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + grid_t = image_grid_thw[image_idx][0] + grid_hs = image_grid_thw[:, 1] + grid_ws = image_grid_thw[:, 2] + t_index = torch.arange(grid_t) * position_id_per_seconds + llm_pos_ids = get_llm_pos_ids_for_vision( + st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws + ) + image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2) + llm_pos_ids_list.append(llm_pos_ids) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + eos_len = 1 + llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + st += text_len + bos_len + image_len + eos_len + image_idx += 1 + remain_images -= 1 + elif ( + min_ed == ed_vision_start + and input_ids[ed_vision_start + 1] == video_token_id + and not use_audio_in_video + ): + text_len = min_ed - st + if text_len != 0: + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + bos_len = 1 + llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + grid_t = video_grid_thw[video_idx][0] + grid_hs = video_grid_thw[:, 1] + grid_ws = video_grid_thw[:, 2] + t_index = torch.arange(grid_t) * float(second_per_grids[video_idx].item()) * position_id_per_seconds + llm_pos_ids = get_llm_pos_ids_for_vision( + st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws + ) + video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2) + llm_pos_ids_list.append(llm_pos_ids) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + eos_len = 1 + llm_pos_ids_list.append(torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + st += text_len + bos_len + video_len + eos_len + video_idx += 1 + remain_videos -= 1 + elif min_ed == ed_vision_start and ed_vision_start + 1 == ed_audio_start and use_audio_in_video: + text_len = min_ed - st + if text_len != 0: + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + bos_len = 1 + bos_block = torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx + llm_pos_ids_list.append(bos_block) + llm_pos_ids_list.append(bos_block) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + _, audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx]) + audio_llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx + grid_t = video_grid_thw[video_idx][0] + grid_hs = video_grid_thw[:, 1] + grid_ws = video_grid_thw[:, 2] + t_index = torch.arange(grid_t) * float(second_per_grids[video_idx].item()) * position_id_per_seconds + video_llm_pos_ids = get_llm_pos_ids_for_vision( + st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws + ) + video_data_index, audio_data_index = 0, 0 + while video_data_index < video_llm_pos_ids.shape[-1] and audio_data_index < audio_llm_pos_ids.shape[-1]: + if video_llm_pos_ids[0][video_data_index] <= audio_llm_pos_ids[0][audio_data_index]: + llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_data_index + 1]) + video_data_index += 1 + else: + llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_data_index + 1]) + audio_data_index += 1 + if video_data_index < video_llm_pos_ids.shape[-1]: + llm_pos_ids_list.append(video_llm_pos_ids[:, video_data_index : video_llm_pos_ids.shape[-1]]) + if audio_data_index < audio_llm_pos_ids.shape[-1]: + llm_pos_ids_list.append(audio_llm_pos_ids[:, audio_data_index : audio_llm_pos_ids.shape[-1]]) + video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2) + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + eos_len = 1 + eos_block = torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx + llm_pos_ids_list.append(eos_block) + llm_pos_ids_list.append(eos_block) + st += text_len + bos_len * 2 + audio_len + video_len + eos_len * 2 # noqa: E501 + audio_idx += 1 + video_idx += 1 + remain_videos -= 1 + remain_audios -= 1 + + if st < len(input_tokens): + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 + text_len = len(input_tokens) - st + llm_pos_ids_list.append(torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) + + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1) + if llm_positions.shape[1] != seq_len: + raise RuntimeError("Position ids length mismatch with input ids length") + + mrope_position_delta = llm_positions.max() + 1 - seq_len + return llm_positions, mrope_position_delta diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index 8e6977950..fd64b8b2e 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -27,6 +27,27 @@ _OMNI_MODELS = { "Qwen2_5OmniToken2WavModel", ), "Qwen2ForCausalLM_old": ("qwen2_5_omni", "qwen2_old", "Qwen2ForCausalLM"), # need to discuss + # Qwen3 Omni MoE models + "Qwen3OmniMoeForConditionalGeneration": ( + "qwen3_omni", + "qwen3_omni", + "Qwen3OmniMoeForConditionalGeneration", + ), + "Qwen3OmniMoeThinkerForConditionalGeneration": ( + "qwen3_omni", + "qwen3_omni_moe_thinker", + "Qwen3OmniMoeThinkerForConditionalGeneration", + ), + "Qwen3OmniMoeTalkerForConditionalGeneration": ( + "qwen3_omni", + "qwen3_omni_moe_talker", + "Qwen3OmniMoeTalkerForConditionalGeneration", + ), + "Qwen3OmniMoeCode2Wav": ( + "qwen3_omni", + "qwen3_omni_code2wav", + "Qwen3OmniMoeCode2Wav", + ), } _VLLM_OMNI_MODELS = { diff --git a/vllm_omni/model_executor/models/utils.py b/vllm_omni/model_executor/models/utils.py index f3854ad4e..7f73067b3 100644 --- a/vllm_omni/model_executor/models/utils.py +++ b/vllm_omni/model_executor/models/utils.py @@ -28,3 +28,12 @@ def split_list_into_ranges(lst: torch.Tensor, interval: int) -> list[list[int]]: ranges[index].append(num) return ranges + + +def safe_tensor_reshape(tensor: torch.Tensor, shape: tuple) -> torch.Tensor: + """ + Reshape a tensor safely. + """ + if tensor is None: + return None + return tensor.reshape(shape) diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml index e4f6b8e7b..c378a5b48 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -30,7 +30,7 @@ stage_args: runtime: process: true devices: "1" - max_batch_size: 3 + max_batch_size: 1 engine_args: model_stage: talker model_arch: Qwen2_5OmniForConditionalGeneration diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml new file mode 100644 index 000000000..8584b4c6d --- /dev/null +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml @@ -0,0 +1,94 @@ +# Stage config for running Qwen3-Omni-MoE with 3-stage architecture +# Stage 0: Thinker (multimodal understanding + text generation) +# Stage 1: Talker (text embeddings → 8-layer RVQ codec codes) +# Stage 2: Code2Wav (8-layer RVQ codes → audio waveform) + +stage_args: + - stage_id: 0 + runtime: + devices: "0" + max_batch_size: 1 + engine_args: + model_stage: thinker + model_arch: Qwen3OmniMoeForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler + gpu_memory_utilization: 0.8 + enforce_eager: true + trust_remote_code: true + engine_output_type: latent # Output hidden states for talker + distributed_executor_backend: "mp" + enable_prefix_caching: false + hf_config_name: thinker_config + # tensor_parallel_size: 2 # Commented out - using single GPU for thinker + final_output: true + final_output_type: text + is_comprehension: true + default_sampling_params: + temperature: 0.4 + top_p: 0.9 + top_k: 1 + max_tokens: 2048 + seed: 42 + detokenize: True + repetition_penalty: 1.05 + + - stage_id: 1 + runtime: + devices: "1" + max_batch_size: 1 + engine_args: + model_stage: talker + model_arch: Qwen3OmniMoeForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker + scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler + gpu_memory_utilization: 0.5 + enforce_eager: true + trust_remote_code: true + engine_output_type: latent # Output codec codes for code2wav + # tensor_parallel_size: 2 + enable_prefix_caching: false + distributed_executor_backend: "mp" + hf_config_name: talker_config + engine_input_source: [0] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.thinker2talker + # final_output: true + # final_output_type: text + default_sampling_params: + temperature: 0.9 + top_k: 50 + max_tokens: 4096 + seed: 42 + detokenize: False + repetition_penalty: 1.05 + stop_token_ids: [2150] + + - stage_id: 2 + runtime: + devices: "1" + max_batch_size: 1 + engine_args: + model_stage: code2wav + model_arch: Qwen3OmniMoeForConditionalGeneration + worker_cls: vllm_omni.worker.gpu_generation_worker.GPUGenerationWorker + scheduler_cls: vllm_omni.core.sched.generation_scheduler.GenerationScheduler + enforce_eager: true + trust_remote_code: true + enable_prefix_caching: false + engine_output_type: audio # Final output: audio waveform + gpu_memory_utilization: 0.1 + distributed_executor_backend: "mp" + max_num_batched_tokens: 1000000 + hf_config_name: thinker_config + engine_input_source: [1] + custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav + final_output: true + final_output_type: audio + default_sampling_params: + temperature: 0.0 + top_p: 1.0 + top_k: -1 + max_tokens: 65536 + seed: 42 + detokenize: True + repetition_penalty: 1.1 diff --git a/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py new file mode 100644 index 000000000..e2d07aa1b --- /dev/null +++ b/vllm_omni/model_executor/stage_input_processors/qwen3_omni.py @@ -0,0 +1,171 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +# Copyright 2025 The Qwen team. +"""Stage input processor for Qwen3 Omni MoE: Thinker → Talker transition.""" + +from typing import Union + +import torch +from vllm.inputs import TextPrompt + +from vllm_omni.inputs.data import OmniTokensPrompt + + +def _compute_talker_prompt_ids_length(info): + im_start_token_id = 151644 + system_token_id = 8948 + user_token_id = 872 + assistant_token_id = 77091 + + thinker_sequences = torch.tensor(info["thinker_sequences"]).unsqueeze(0).cuda().long() # [1, T] + input_ids = torch.tensor(info["thinker_input_ids"]).unsqueeze(0).cuda().long() # [1, T] + + im_start_indexes = torch.cat( + [ + torch.nonzero(input_ids[0] == im_start_token_id).squeeze(1), + torch.tensor([thinker_sequences.shape[-1]], device=input_ids.device, dtype=input_ids.dtype), + ], + dim=0, + ) + + sum_user_len = 0 + assistant_len = 0 + for i in range(len(im_start_indexes) - 1): + s = int(im_start_indexes[i].item()) + e = int(im_start_indexes[i + 1].item()) + role = int(input_ids[0, s + 1].item()) + if role == system_token_id: + continue + elif role == user_token_id: + sum_user_len += e - s + elif role == assistant_token_id and i == len(im_start_indexes) - 2: + assistant_len += 9 # 3 + 4 + 1 + 1 + else: + pass + + return sum_user_len + assistant_len + + +def thinker2talker( + stage_list, + engine_input_source, + prompt: Union[OmniTokensPrompt, TextPrompt] = None, + requires_multimodal_data: bool = False, +): + """ + Process thinker outputs to create talker inputs. + + Workflow: + 1. Extract thinker's text generation outputs (token IDs + hidden states) + 2. Split hidden states into: prompt embeddings + generated embeddings + 3. Package for talker with additional information + + Args: + stage_list: List of stage objects + engine_input_source: Source stage IDs (typically [0] for thinker) + prompt: Original prompt data + + Returns: + List of OmniTokensPrompt for talker stage + """ + 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}") + + if stage_list[source_stage_id].engine_outputs is None: + raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") + + thinker_outputs = stage_list[source_stage_id].engine_outputs + talker_inputs = [] + + # Process each thinker output + for i, thinker_output in enumerate(thinker_outputs): + output = thinker_output.outputs[0] + thinker_embeddings = output.multimodal_output["0"].float().clone().detach().cuda() + + thinker_hidden_states = output.multimodal_output["24"].float().clone().detach().cuda() + info = { + "thinker_embeddings": thinker_embeddings, + "thinker_hidden_states": thinker_hidden_states, + "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.get("tts_bos_embed").float().clone().detach().cuda(), + "tts_eos_embed": output.multimodal_output.get("tts_eos_embed").float().clone().detach().cuda(), + "tts_pad_embed": output.multimodal_output.get("tts_pad_embed").float().clone().detach().cuda(), + } + talker_inputs.append( + OmniTokensPrompt( + prompt_token_ids=[0] * _compute_talker_prompt_ids_length(info), + additional_information=info, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + + return talker_inputs + + +def talker2code2wav( + stage_list, + engine_input_source, + prompt: Union[OmniTokensPrompt, TextPrompt] = None, + requires_multimodal_data: bool = False, +): + """ + Process talker outputs to create code2wav inputs. + + Workflow: + 1. Extract talker's codec code outputs (8-layer RVQ codes) + 2. Flatten codes for code2wav input + 3. Package for code2wav stage + + Args: + stage_list: List of stage objects + engine_input_source: Source stage IDs (typically [1] for talker) + prompt: Original prompt data + + Returns: + List of OmniTokensPrompt for code2wav stage + """ + 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}") + + if stage_list[source_stage_id].engine_outputs is None: + raise RuntimeError(f"Stage {source_stage_id} has no outputs yet") + + talker_outputs = stage_list[source_stage_id].engine_outputs + code2wav_inputs = [] + + # Process each talker output + for i, talker_output in enumerate(talker_outputs): + output = talker_output.outputs[0] + + # Extract codec codes from talker output + # Expected shape: [8, seq_len] (8-layer RVQ codes) + codec_codes = ( + output.multimodal_output["code_predictor_codes"] + .to(torch.long) + .transpose(0, 1) + .cpu() + .to(torch.long) + .reshape(-1) + .tolist() + ) # 16, seq_len + code2wav_inputs.append( + OmniTokensPrompt( + prompt_token_ids=codec_codes, + multi_modal_data=None, + mm_processor_kwargs=None, + ) + ) + + return code2wav_inputs diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index ea0828cff..a0028d572 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -1,16 +1,24 @@ import sys from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt -from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding as _OriginalMRotaryEmbedding +from vllm.model_executor.layers.rotary_embedding import ( + MRotaryEmbedding as _OriginalMRotaryEmbedding, +) +from vllm.v1.engine import EngineCoreOutput as _OriginalEngineCoreOutput +from vllm.v1.engine import EngineCoreOutputs as _OriginalEngineCoreOutputs from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest from vllm.v1.request import Request as _OriginalRequest -from vllm_omni.engine import OmniEngineCoreRequest +from vllm_omni.engine import OmniEngineCoreOutput, OmniEngineCoreOutputs, OmniEngineCoreRequest from vllm_omni.inputs.data import OmniTokensPrompt from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding from vllm_omni.request import OmniRequest for module_name, module in sys.modules.items(): + if hasattr(module, "EngineCoreOutput") and module.EngineCoreOutput == _OriginalEngineCoreOutput: + module.EngineCoreOutput = OmniEngineCoreOutput + if hasattr(module, "EngineCoreOutputs") and module.EngineCoreOutputs == _OriginalEngineCoreOutputs: + module.EngineCoreOutputs = OmniEngineCoreOutputs if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt: module.TokensPrompt = OmniTokensPrompt if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding: diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 5830a3adc..c78aade3d 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -6,9 +6,6 @@ and also outputs sampled tokens. from __future__ import annotations -from typing import Any - -import numpy as np import torch from vllm.forward_context import BatchDescriptor from vllm.logger import init_logger @@ -25,10 +22,8 @@ from vllm.v1.worker.gpu_model_runner import ( has_kv_transfer_group, set_forward_context, ) -from vllm.v1.worker.ubatch_utils import UBatchSlices from vllm.v1.worker.utils import is_residual_scattered_for_sp -from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload from vllm_omni.outputs import OmniModelRunnerOutput from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner @@ -48,169 +43,9 @@ class GPUARModelRunner(OmniGPUModelRunner): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self.input_ids = self._make_buffer(self.max_num_tokens, dtype=torch.int32) - - def _preprocess( - self, - scheduler_output: SchedulerOutput, - num_scheduled_tokens_np: np.ndarray, - intermediate_tensors: IntermediateTensors | None = None, - ubatch_slices: UBatchSlices | None = None, - num_tokens_after_padding: torch.Tensor | None = None, - ) -> tuple[ - int, - int, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor, - IntermediateTensors | None, - dict[str, Any], - dict[str, dict] | None, - ]: - num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens - if ubatch_slices: - assert num_tokens_after_padding is not None - num_input_tokens = int(num_tokens_after_padding[0].item() * 2) - self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens) - elif ubatch_slices is None: - num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens) - num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens) - num_input_tokens += num_pad - - # _prepare_inputs may reorder the batch, so we must gather multi - # modal outputs after that to ensure the correct order - per_req_additional_information: dict[str, dict] | None = None - if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder: - # Run the multimodal encoder if any. - self._execute_mm_encoder(scheduler_output) - mm_embeds = self._gather_mm_embeddings(scheduler_output) - - # NOTE(woosuk): To unify token ids and soft tokens (vision - # embeddings), we always use embeddings (rather than token ids) - # as input to the multimodal model, even when the input is text. - inputs_embeds_scheduled = self.model.get_input_embeddings( - input_ids=self.input_ids.gpu[:num_scheduled_tokens], - multimodal_embeddings=mm_embeds or None, - ) - - # TODO(woosuk): Avoid the copy. Optimize. - self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) - - # Reset per-step additional information collector (deprecated concat path) - if hasattr(self, "_forward_additional_information"): - self._forward_additional_information = None - # New: per-request additional information for this step - per_req_additional_information = {} - - # Overlay custom prompt_embeds per request for the prompt portion; - # collect additional_information (tensor/list) for prefill portion only - for req_index, req_id in enumerate(self.input_batch.req_ids): - req_state = self.requests[req_id] - pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) - addi_cpu = getattr(req_state, "additional_information_cpu", None) - num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index]) - prompt_len = len(req_state.prompt_token_ids) - prompt_remaining = max(0, prompt_len - num_computed_tokens) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - overlay_len = min(sched_tokens, prompt_remaining) - if overlay_len <= 0: - continue - if pe_cpu is not None: - src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to( - dtype=self.dtype, device=self.device, non_blocking=True - ) - start_offset = int(self.query_start_loc.cpu[req_index]) - self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src) - # Build per-request additional information (no cross-request concat) - if addi_cpu is not None and isinstance(addi_cpu, dict): - req_info: dict[str, object] = {} - for k, v in addi_cpu.items(): - if isinstance(v, torch.Tensor): - # For prefill tokens, pass only the scheduled slice; - # for decode or no scheduled tokens, pass whole tensor - if overlay_len > 0: - try: - seg = ( - v[num_computed_tokens : num_computed_tokens + overlay_len] - .detach() - .to("cpu") - .contiguous() - ) - except Exception: - seg = v.detach().to("cpu").contiguous() - req_info[k] = seg - else: - req_info[k] = v.detach().to("cpu").contiguous() - elif isinstance(v, list): - req_info[k] = v - else: - req_info[k] = v - per_req_additional_information[req_id] = req_info - - input_ids = self.input_ids.gpu[:num_input_tokens] - inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] - model_kwargs = { - **self._init_model_kwargs(num_scheduled_tokens), - **self._extract_mm_kwargs(scheduler_output), - } - elif self.enable_prompt_embeds and get_pp_group().is_first_rank: - # Get the input embeddings for the tokens that are not input embeds, - # then put them into the appropriate positions. - # TODO(qthequartermasterman): Since even when prompt embeds are - # enabled, (a) not all requests will use prompt embeds, and (b) - # after the initial prompt is processed, the rest of the generated - # tokens will be token ids, it is not desirable to have the - # embedding layer outside of the CUDA graph all the time. The v0 - # engine avoids this by "double compiling" the CUDA graph, once - # with input_ids and again with inputs_embeds, for all num_tokens. - # If a batch only has token ids, then including the embedding layer - # in the CUDA graph will be more performant (like in the else case - # below). - token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1) - # Some tokens ids may need to become embeds - if token_ids_idx.numel() > 0: - token_ids = self.input_ids.gpu[token_ids_idx] - tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids) - self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds - - inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] - model_kwargs = self._init_model_kwargs(num_input_tokens) - input_ids = None - else: - # For text-only models, we use token ids as input. - # While it is possible to use embeddings as input just like the - # multimodal models, it is not desirable for performance since - # then the embedding layer is not included in the CUDA graph. - input_ids = self.input_ids.gpu[:num_input_tokens] - inputs_embeds = None - model_kwargs = self._init_model_kwargs(num_input_tokens) - if self.uses_mrope: - positions = self.mrope_positions.gpu[:, :num_input_tokens] - else: - positions = self.positions.gpu[:num_input_tokens] - - if get_pp_group().is_first_rank: - intermediate_tensors = None - else: - intermediate_tensors = self.sync_and_slice_intermediate_tensors( - num_input_tokens, intermediate_tensors, True - ) - - if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: - encoder_inputs = self._extract_encoder_inputs(scheduler_output) - model_kwargs.update(encoder_inputs) - - return ( - num_scheduled_tokens, - num_input_tokens, - num_tokens_after_padding, - input_ids, - inputs_embeds, - positions, - intermediate_tensors, - model_kwargs, - per_req_additional_information, - ) + # each model stage has their own hidden size + self.hidden_size = self.model_config.hf_text_config.hidden_size + self.inputs_embeds = self._make_buffer(self.max_num_tokens, self.hidden_size, dtype=self.dtype, numpy=False) @torch.inference_mode() def execute_model( @@ -223,58 +58,9 @@ class GPUARModelRunner(OmniGPUModelRunner): # Update persistent batch states. self._update_states(scheduler_output) - # Decode per-request prompt_embeds / additional_hidden_states payloads + # Decode per-request prompt_embeds / additional_information payloads # (if present) into CPU tensors - try: - new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) - if new_reqs: - for nr in new_reqs: - req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) - if req_id is None: - continue - # prompt_embeds - payload_pe = getattr(nr, "prompt_embeds", None) - if payload_pe is not None: - if isinstance(payload_pe, torch.Tensor): - pe_cpu = payload_pe.detach().to("cpu").contiguous() - elif isinstance(payload_pe, PromptEmbedsPayload): - dt = np.dtype(getattr(payload_pe, "dtype", "float32")) - arr = np.frombuffer(payload_pe.data, dtype=dt) - arr = arr.reshape(payload_pe.shape) - pe_cpu = torch.from_numpy(arr) - else: - pe_cpu = None - if pe_cpu is not None and req_id in self.requests: - setattr( - self.requests[req_id], - "prompt_embeds_cpu", - pe_cpu, - ) - # additional_information - payload_info = getattr(nr, "additional_information", None) - if payload_info is not None: - info_dict = {} - if isinstance(payload_info, dict): - # Already decoded - info_dict = payload_info - elif isinstance(payload_info, AdditionalInformationPayload): - for k, entry in payload_info.entries.items(): - if entry.tensor_data is not None: - dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) - arr = np.frombuffer(entry.tensor_data, dtype=dt) - arr = arr.reshape(entry.tensor_shape) - info_dict[k] = torch.from_numpy(arr) - else: - info_dict[k] = entry.list_data - if info_dict and req_id in self.requests: - setattr( - self.requests[req_id], - "additional_information_cpu", - info_dict, - ) - except Exception as e: - logger.error(f"Error decoding prompt_embeds / additional_information: {e}") - pass + self._decode_and_store_request_payloads(scheduler_output) if not scheduler_output.total_num_scheduled_tokens: if not has_kv_transfer_group(): @@ -344,29 +130,7 @@ class GPUARModelRunner(OmniGPUModelRunner): record_function_or_nullcontext("Forward"), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, ): - model_kwargs_extra = {} - # Pass per-request additional information map for this step (no concat) - if per_req_additional_information: - model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information - # Always pass per-request runtime additional_information (persisted in request state) - try: - per_req_runtime_info = [] - for req_id in self.input_batch.req_ids: - req_state = self.requests.get(req_id) - info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - per_req_runtime_info.append(info if isinstance(info, dict) else {}) - model_kwargs_extra["runtime_additional_information"] = per_req_runtime_info - model_kwargs_extra["request_ids"] = self.input_batch.req_ids - # Pass each request's token span within the flattened sequence for this step, - # enabling the model to map decode/prefill by request - req_token_spans = [] - for req_index in range(len(self.input_batch.req_ids)): - start_offset = int(self.query_start_loc.cpu[req_index]) - sched_tokens = int(num_scheduled_tokens_np[req_index]) - req_token_spans.append((start_offset, start_offset + sched_tokens)) - model_kwargs_extra["request_token_spans"] = req_token_spans - except Exception: - pass + model_kwargs_extra = self._build_model_kwargs_extra(per_req_additional_information, num_scheduled_tokens_np) model_output = self.model( input_ids=input_ids, positions=positions, @@ -388,37 +152,15 @@ class GPUARModelRunner(OmniGPUModelRunner): hidden_states = model_output aux_hidden_states = None - hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) + # hidden_states, multimodal_outputs = self.extract_multimodal_outputs( + # hidden_states + # ) + multimodal_outputs = model_output.multimodal_outputs + hidden_states = model_output.text_hidden_states # The model side may return per-request additional_information updates (model-agnostic channel). # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order; # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode. - try: - if isinstance(multimodal_outputs, dict) and ( - "additional_information_update" in multimodal_outputs - or "additional_information_update_by_req_id" in multimodal_outputs - ): - # Option A: list[dict] in batch order - updates_list = multimodal_outputs.get("additional_information_update") - if isinstance(updates_list, list): - for idx, upd in enumerate(updates_list): - if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids): - continue - req_id = self.input_batch.req_ids[idx] - self._merge_additional_information_update(req_id, upd) - # Option B: dict[str, dict] keyed by req_id - updates_map = multimodal_outputs.get("additional_information_update_by_req_id") - if isinstance(updates_map, dict): - for req_id, upd in updates_map.items(): - if not isinstance(upd, dict): - continue - if req_id not in self.requests: - continue - self._merge_additional_information_update(req_id, upd) - except Exception as e: - logger.error( - f"Error merging for requests:{self.input_batch.req_ids} additional \ - information update: {e}, with the multimodal_outputs as {multimodal_outputs}" - ) + self._process_additional_information_updates(multimodal_outputs) if not self.broadcast_pp_output: # Common case. if not get_pp_group().is_last_rank: @@ -534,12 +276,42 @@ class GPUARModelRunner(OmniGPUModelRunner): # Convert to per-request tensors on CPU hidden_states_cpu = hidden_states.detach().to("cpu").contiguous() + # pooler_output: list[torch.Tensor | None] = [] pooler_output: list[torch.Tensor | None] = [] prev_logits_index = 0 - for logits_index in logits_indices: - pooler_output.append(hidden_states_cpu[prev_logits_index : logits_index + 1]) + for rid, logits_index in zip(req_ids_output_copy, logits_indices): + # Base payload: hidden slice for this request in this iteration + hidden_slice = hidden_states_cpu[prev_logits_index : logits_index + 1] + payload: dict[str, object] = {"hidden": hidden_slice} + # Merge multimodal_outputs if present + if isinstance(multimodal_outputs, dict) and multimodal_outputs: + mm_payload: dict[str, object] = {} + for k, v in multimodal_outputs.items(): + try: + # Case 1: tensor aligned on token dimension + if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: + mm_payload[k] = v.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous() + # Case 2: nested dict of tensors aligned on token dimension (e.g., selected_hidden_layers) + 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")[prev_logits_index : logits_index + 1].contiguous() + ) + if sub_dict: + mm_payload[k] = sub_dict + elif isinstance(v, list): + element: torch.Tensor = v[0] + multimodal_outputs[k] = v[1:] if len(v) > 1 else v + mm_payload[k] = element + except Exception as e: + # Best-effort; skip malformed entries + logger.error(f"Error in merge multimodal outputs: {e}") + if mm_payload: + payload.update(mm_payload) + pooler_output.append(payload) # type: ignore[arg-type] prev_logits_index = logits_index + 1 - output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -583,3 +355,5 @@ class GPUARModelRunner(OmniGPUModelRunner): else: merged[k] = v setattr(req_state, "additional_information_cpu", merged) + + # ===== Helper functions extracted for clarity and reuse ===== diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py new file mode 100644 index 000000000..d18b7e956 --- /dev/null +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -0,0 +1,582 @@ +"""Code2Wav GPU Model Runner for vLLM-omni. + +Handles direct conversion from codec codes to audio waveforms for Qwen3 Omni MoE Code2Wav. +This is a non-autoregressive model that doesn't require sampling or logits computation. +""" + +from __future__ import annotations + +import gc +import logging +from typing import Any + +import numpy as np +import torch +from vllm.config import CUDAGraphMode +from vllm.forward_context import BatchDescriptor +from vllm.multimodal.inputs import MultiModalKwargs +from vllm.v1.attention.backends.utils import CommonAttentionMetadata +from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.spec_decode.eagle import EagleProposer +from vllm.v1.utils import record_function_or_nullcontext +from vllm.v1.worker.gpu_model_runner import ( + EMPTY_MODEL_RUNNER_OUTPUT, + AsyncGPUModelRunnerOutput, + IntermediateTensors, + PerLayerAttnMetadata, + get_pp_group, + set_forward_context, +) +from vllm.v1.worker.ubatch_utils import UBatchSlices +from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs + +from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner + +logger = logging.getLogger(__name__) + + +class GPUGenerationModelRunner(OmniGPUModelRunner): + """Generation model runner for vLLM-omni (non-autoregressive). + + - Reuses GPUModelRunner preparation, multimodal handling, and TP/PP/DP glue. + - Does not compute logits or perform token sampling. + - Executes generation process and returns tensors via `pooler_output`. + """ + + def _preprocess( + self, + scheduler_output: SchedulerOutput, + intermediate_tensors: IntermediateTensors | None = None, + ubatch_slices: UBatchSlices | None = None, + num_tokens_after_padding: torch.Tensor | None = None, + ) -> tuple[ + int, + int, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor | None, + torch.Tensor, + IntermediateTensors | None, + dict[str, Any], + ]: + # Input token count for this iteration (not used by diffusion, but + # retained to keep DP padding/ordering consistent) + num_input_tokens = scheduler_output.total_num_scheduled_tokens + num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens) + num_input_tokens += num_pad + + # _prepare_inputs may reorder the batch, so we must gather multi + # modal outputs after that to ensure the correct order + if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder: + # Run the multimodal encoder if any. + self._execute_mm_encoder(scheduler_output) + mm_embeds = self._gather_mm_embeddings(scheduler_output) + + # NOTE(woosuk): To unify token ids and soft tokens (vision + # embeddings), we always use embeddings (rather than token ids) + # as input to the multimodal model, even when the input is text. + inputs_embeds_scheduled = self.model.get_input_embeddings( + input_ids=self.input_ids.gpu[:num_input_tokens], + multimodal_embeddings=mm_embeds or None, + ) + + # TODO(woosuk): Avoid the copy. Optimize. + self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled) + + input_ids = self.input_ids.gpu[:num_input_tokens] + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = { + **self._init_model_kwargs(num_input_tokens), + **self._extract_mm_kwargs(scheduler_output), + } + elif self.enable_prompt_embeds and get_pp_group().is_first_rank: + # Get the input embeddings for the tokens that are not input embeds, + # then put them into the appropriate positions. + # TODO(qthequartermasterman): Since even when prompt embeds are + # enabled, (a) not all requests will use prompt embeds, and (b) + # after the initial prompt is processed, the rest of the generated + # tokens will be token ids, it is not desirable to have the + # embedding layer outside of the CUDA graph all the time. The v0 + # engine avoids this by "double compiling" the CUDA graph, once + # with input_ids and again with inputs_embeds, for all num_tokens. + # If a batch only has token ids, then including the embedding layer + # in the CUDA graph will be more performant (like in the else case + # below). + token_ids_idx = self.is_token_ids.gpu[:num_input_tokens].nonzero(as_tuple=False).squeeze(1) + # Some tokens ids may need to become embeds + if token_ids_idx.numel() > 0: + token_ids = self.input_ids.gpu[token_ids_idx] + tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids) + self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds + + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = self._init_model_kwargs(num_input_tokens) + input_ids = self.input_ids.gpu[:num_input_tokens] + else: + # For text-only models, we use token ids as input. + # While it is possible to use embeddings as input just like the + # multimodal models, it is not desirable for performance since + # then the embedding layer is not included in the CUDA graph. + input_ids = self.input_ids.gpu[:num_input_tokens] + inputs_embeds = None + model_kwargs = self._init_model_kwargs(num_input_tokens) + if self.uses_mrope: + positions = self.mrope_positions.gpu[:, :num_input_tokens] + else: + positions = self.positions.gpu[:num_input_tokens] + + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + intermediate_tensors = self.sync_and_slice_intermediate_tensors( + num_input_tokens, intermediate_tensors, True + ) + + if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: + encoder_inputs = self._extract_encoder_inputs(scheduler_output) + model_kwargs.update(encoder_inputs) + + return ( + num_input_tokens, + num_input_tokens, + num_tokens_after_padding, + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ) + + @torch.inference_mode() + def execute_model( + self, + scheduler_output: SchedulerOutput, + intermediate_tensors: IntermediateTensors | None = None, + ) -> OmniModelRunnerOutput | IntermediateTensors: + with record_function_or_nullcontext("Preprocess"): + with self.synchronize_input_prep(): + # Update persistent batch states. + self._update_states(scheduler_output) + if not scheduler_output.total_num_scheduled_tokens: + return EMPTY_MODEL_RUNNER_OUTPUT + + # Prepare the decoder inputs. + ( + attn_metadata, + logits_indices, + spec_decode_metadata, + num_scheduled_tokens_np, + spec_decode_common_attn_metadata, + max_query_len, + ubatch_slices, + num_tokens_after_padding, + ) = self._prepare_inputs(scheduler_output) + + # Input token count for this iteration (not used by diffusion, but + # retained to keep DP padding/ordering consistent) + ( + num_scheduled_tokens, + num_input_tokens, + num_tokens_across_dp, + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ) = self._preprocess( + scheduler_output, + intermediate_tensors, + ubatch_slices, + num_tokens_after_padding, + ) + + cudagraph_runtime_mode = CUDAGraphMode.NONE + # Run the model. + # Use persistent buffers for CUDA graphs. + with ( + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_input_tokens, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=None, + ubatch_slices=ubatch_slices, + ), + record_function_or_nullcontext("Forward"), + self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, + ): + outputs = self._run_generation_model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + multimodal_kwargs=model_kwargs, + logits_indices=logits_indices, + ) + + _, multimodal_outputs = self.extract_multimodal_outputs(outputs) + # Ensure one tensor per request, map to CPU for output struct + pooler_output: list[torch.Tensor | None] = [] + if isinstance(multimodal_outputs, torch.Tensor): + # If model returned a single stacked tensor, split by requests + assert multimodal_outputs.shape[0] == self.input_batch.num_reqs + for i in range(self.input_batch.num_reqs): + pooler_output.append(multimodal_outputs[i].detach().to("cpu").contiguous()) + elif isinstance(multimodal_outputs, list): + for out in multimodal_outputs: + pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None) + elif isinstance(multimodal_outputs, dict): + for out in multimodal_outputs.values(): + pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None) + else: + raise RuntimeError("Unsupported diffusion output type") + + output = OmniModelRunnerOutput( + req_ids=self.input_batch.req_ids, + req_id_to_index=self.input_batch.req_id_to_index, + sampled_token_ids=[], + logprobs=None, + prompt_logprobs_dict={}, + pooler_output=pooler_output, + kv_connector_output=kv_connector_output, + num_nans_in_logits={}, + ) + + if not self.use_async_scheduling: + return output + + return AsyncGPUModelRunnerOutput( + model_runner_output=output, + sampled_token_ids=[], + invalid_req_indices=[], + async_output_copy_stream=self.async_output_copy_stream, + ) + + def _run_generation_model( + self, + *, + input_ids: torch.Tensor, + positions: torch.Tensor, + intermediate_tensors: IntermediateTensors | None, + inputs_embeds: torch.Tensor | None, + multimodal_kwargs: dict, + logits_indices: torch.Tensor, + ) -> torch.Tensor | list[torch.Tensor]: + """Run generation from codec codes to waveforms. + + Args: + scheduler_output: Contains codec codes in input_ids or additional info + intermediate_tensors: PP intermediate tensors if applicable + + Returns: + Audio waveforms: [batch, 1, waveform_len] or list of tensors + """ + # Keep inputs identical to AR runner + kwargs = dict( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device), + sampling_metadata=self.input_batch.sampling_metadata, + logits_index=logits_indices, + sampler=self.sampler, + ) + + if hasattr(self.model, "forward"): + return self.model.forward(**kwargs) + + raise RuntimeError( + "The loaded model does not expose diffusion interfaces 'sample', " + "'forward', or 'diffuse'. Please implement one of them or adapt the runner." + ) + + @torch.inference_mode() + def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None: + logger.warning("Dummy sampler run is not implemented for generation model") + return None + + @torch.inference_mode() + def _dummy_run( + self, + num_tokens: int, + cudagraph_runtime_mode: CUDAGraphMode | None = None, + force_attention: bool = False, + uniform_decode: bool = False, + skip_eplb: bool = False, + is_profile: bool = False, + create_mixed_batch: bool = False, + remove_lora: bool = True, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Run a dummy forward pass to warm up/profile run or capture the + CUDA graph for the model. + + Args: + num_tokens: Number of tokens to run the dummy forward pass. + cudagraph_runtime_mode: used to control the behavior. + - if not set will determine the cudagraph mode based on using + the self.cudagraph_dispatcher. + - CUDAGraphMode.NONE: No cudagraph, for warm up and profile run + - CUDAGraphMode.PIECEWISE: Piecewise cudagraph. + - CUDAGraphMode.FULL: Full cudagraph, attention metadata is + needed. + force_attention: If True, always create attention metadata. Used to + warm up attention backend when mode is NONE. + uniform_decode: If True, the batch is a uniform decode batch. + skip_eplb: If True, skip EPLB state update. + is_profile: If True, this is a profile run. + create_mixed_batch: If True, create a mixed batch with both decode + (1 token) and prefill (multiple tokens) requests. + remove_lora: If False, dummy LoRAs are not destroyed after the run + """ + assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in { + CUDAGraphMode.NONE, + CUDAGraphMode.PIECEWISE, + CUDAGraphMode.FULL, + } + + # If cudagraph_mode.decode_mode() == FULL and + # cudagraph_mode.separate_routine(). This means that we are using + # different graphs and/or modes for mixed prefill-decode batches vs. + # uniform decode batches. A uniform decode batch means that all + # requests have identical query length, except a potential virtual + # request (shorter) in the batch account for padding. + # Uniform decode batch could either be common pure decode, where + # max_query_len == 1, or speculative decode, where + # max_query_len == 1 + num_spec_decode_tokens. + + # When setting max_query_len = 1, we switch to and capture the optimized + # routine of FA2 for pure decode, i.e., Flashdecode + an optimization + # for GQA/MQA. + max_query_len = self.uniform_decode_query_len if uniform_decode else num_tokens + + # Set num_scheduled_tokens based on num_tokens and max_num_seqs + # for dummy run with LoRA so that the num_reqs collectively + # has num_tokens in total. + assert num_tokens <= self.scheduler_config.max_num_batched_tokens + max_num_reqs = self.scheduler_config.max_num_seqs + + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs + num_scheduled_tokens_list[-1] += num_tokens % num_reqs + + assert sum(num_scheduled_tokens_list) == num_tokens + assert len(num_scheduled_tokens_list) == num_reqs + num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) + + num_tokens_after_padding = None + + # If we failed to microbatch, currently need to resynchronize + # TODO(lucas,sage): we should be able to avoid this second sync by + # refactoring `get_dp_padding_ubatch` and `get_dp_padding` into + # a single `coordinate_batch_across_dp` function. + if num_tokens_after_padding is None: + num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) + num_tokens_after_padding = num_tokens + num_pad + else: + num_tokens_across_dp = num_tokens_after_padding + num_tokens_after_padding = int(num_tokens_after_padding[0].item()) + + attn_metadata: PerLayerAttnMetadata | None = None + + # If force_attention is True, we always capture attention. Otherwise, + # it only happens for cudagraph_runtime_mode=FULL. + if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: + attn_metadata = {} + + seq_lens = max_query_len + self.seq_lens.np[:num_reqs] = seq_lens + self.seq_lens.np[num_reqs:] = 0 + self.seq_lens.copy_to_gpu() + + cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.copy_to_gpu() + + for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups): + common_attn_metadata = CommonAttentionMetadata( + query_start_loc=self.query_start_loc.gpu[: num_reqs + 1], + query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1], + seq_lens=self.seq_lens.gpu[:num_reqs], + seq_lens_cpu=self.seq_lens.cpu[:num_reqs], + num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs], # noqa: E501 + num_reqs=num_reqs, + num_actual_tokens=num_tokens, + max_query_len=max_query_len, + max_seq_len=self.max_model_len, + block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs), + slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens], + causal=True, + ) + for attn_group in self.attn_groups[kv_cache_group_id]: + assert type(attn_metadata) is dict + attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture( + common_attn_metadata + ) # noqa: E501 + for layer_name in attn_group.layer_names: + attn_metadata[layer_name] = attn_metadata_i + + with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora): + model_kwargs = self._init_model_kwargs(num_tokens) + if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: + input_ids = None + inputs_embeds = self.inputs_embeds.gpu[:num_tokens] + model_kwargs = { + **model_kwargs, + **self._dummy_mm_kwargs(num_reqs), + } + elif self.enable_prompt_embeds: + input_ids = None + inputs_embeds = self.inputs_embeds.gpu[:num_tokens] + model_kwargs = self._init_model_kwargs(num_tokens) + else: + input_ids = self.input_ids.gpu[:num_tokens] + inputs_embeds = None + + if self.uses_mrope: + positions = self.mrope_positions.gpu[:, :num_tokens] + else: + positions = self.positions.gpu[:num_tokens] + + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + if self.intermediate_tensors is None: + self.intermediate_tensors = self.model.make_empty_intermediate_tensors( + batch_size=self.max_num_tokens, + dtype=self.model_config.dtype, + device=self.device, + ) + + intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens, None, False) + + # filter out the valid batch descriptor + _cg_mode, batch_descriptor = ( + self.cudagraph_dispatcher.dispatch( + BatchDescriptor( + num_tokens=num_tokens_after_padding, + uniform_decode=uniform_decode, + ) + ) + if not is_profile + else (CUDAGraphMode.NONE, None) + ) + if cudagraph_runtime_mode is not None: + # we allow forcing NONE when the dispatcher disagrees to support + # warm ups for cudagraph capture + assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, ( + f"Cudagraph runtime mode mismatch at dummy_run. " + f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}." + ) + else: + cudagraph_runtime_mode = _cg_mode + + with ( + self.maybe_randomize_inputs(input_ids), + set_forward_context( + attn_metadata, + self.vllm_config, + num_tokens=num_tokens_after_padding, + num_tokens_across_dp=num_tokens_across_dp, + cudagraph_runtime_mode=cudagraph_runtime_mode, + batch_descriptor=batch_descriptor, + ), + ): + outputs = self.model( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + ) + + if self.use_aux_hidden_state_outputs: + hidden_states, _ = outputs + else: + hidden_states = outputs + + if self.speculative_config and self.speculative_config.use_eagle(): + assert isinstance(self.drafter, EagleProposer) + self.drafter.dummy_run(num_tokens) + + # This is necessary to avoid blocking DP. + # For dummy runs, we typically skip EPLB since we don't have any real + # requests to process. + # However, in DP settings, there may be cases when some DP ranks do + # not have any requests to process, so they're executing dummy batches. + # In such cases, we still have to trigger EPLB to make sure + # ranks execute the rearrangement in synchronization. + + # logit_indices = np.cumsum(num_scheduled_tokens) - 1 + hidden_states, _ = self.extract_multimodal_outputs(hidden_states) + return hidden_states, None + + def profile_run(self) -> None: + # Profile with multimodal encoder & encoder cache. + if self.supports_mm_inputs: + if self.model_config.multimodal_config.skip_mm_profiling: + logger.info("Skipping memory profiling for multimodal encoder and encoder cache.") + else: + mm_budget = self.mm_budget + assert mm_budget is not None + + if (encoder_budget := mm_budget.get_encoder_budget()) > 0: + # NOTE: Currently model is profiled with a single non-text + # modality with the max possible input tokens even when + # it supports multiple. + dummy_modality = mm_budget.get_modality_with_max_tokens() + max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[dummy_modality] + + logger.info( + "Encoder cache will be initialized with a budget " + "of %s tokens, and profiled with %s %s items of the maximum feature size.", + encoder_budget, + max_mm_items_per_batch, + dummy_modality, + ) + + # Create dummy batch of multimodal inputs. + batched_dummy_mm_inputs = self._get_mm_dummy_batch( + dummy_modality, + max_mm_items_per_batch, + ) + + # Run multimodal encoder. + dummy_encoder_outputs = self.model.get_multimodal_embeddings(**batched_dummy_mm_inputs) + + sanity_check_mm_encoder_outputs( + dummy_encoder_outputs, + expected_num_items=max_mm_items_per_batch, + ) + + # NOTE: This happens when encoder cache needs to store + # the embeddings that encoder outputs are scattered onto. + # In this case we create dummy embeddings of size + # (encode_budget, hidden_size) and scatter encoder + # output into it. + encoder_output_shape = dummy_encoder_outputs[0].shape + if encoder_output_shape[0] < encoder_budget: + expanded_outputs = [] + for output in dummy_encoder_outputs: + expanded = output.new_zeros((encoder_budget, encoder_output_shape[-1])) + num_tokens = output.shape[0] + expanded[:num_tokens].copy_(output) + expanded_outputs.append(expanded) + + dummy_encoder_outputs = expanded_outputs + + # Cache the dummy encoder outputs. + self.encoder_cache["tmp"] = dict(enumerate(dummy_encoder_outputs)) + + # Add `is_profile` here to pre-allocate communication buffers + hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True) + if get_pp_group().is_last_rank: + pass + self._sync_device() + del hidden_states + self.encoder_cache.clear() + gc.collect() diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py new file mode 100644 index 000000000..49dc10c7b --- /dev/null +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -0,0 +1,84 @@ +import gc +import os + +import torch +from vllm.model_executor import set_random_seed +from vllm.platforms import current_platform +from vllm.utils import GiB_bytes, MemorySnapshot +from vllm.v1.utils import report_usage_stats +from vllm.v1.worker.gpu_worker import Worker as GPUWorker +from vllm.v1.worker.gpu_worker import init_worker_distributed_environment + +from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner + + +class GPUGenerationWorker(GPUWorker): + """GPU Worker for Generation model (non-autoregressive waveform generation). + + Usage in stage config: + worker_cls: "vllm_omni.worker.gpu_generation_model_runner.GPUGenerationModelRunner" + """ + + def init_device(self): + """Initialize CUDA device and distributed environment.""" + if self.device_config.device.type == "cuda": + # torch.distributed.all_reduce does not free the input tensor until + # the synchronization point. This causes the memory usage to grow + # as the number of all_reduce calls increases. This env var disables + # this behavior. + # Related issue: + # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 + os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1" + + # This env var set by Ray causes exceptions with graph building. + os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + self.device = torch.device(f"cuda:{self.local_rank}") + current_platform.set_device(self.device) + + current_platform.check_if_supports_dtype(self.model_config.dtype) + + # Initialize the distributed environment BEFORE taking + # memory snapshot + # This ensures NCCL buffers are allocated before we measure + # available memory + init_worker_distributed_environment( + self.vllm_config, + self.rank, + self.distributed_init_method, + self.local_rank, + current_platform.dist_backend, + ) + + # Set random seed. + set_random_seed(self.model_config.seed) + + # Now take memory snapshot after NCCL is initialized + gc.collect() + torch.cuda.empty_cache() + + # take current memory snapshot + self.init_snapshot = MemorySnapshot() + self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization + if self.init_snapshot.free_memory < self.requested_memory: + + def GiB(b): + return round(b / GiB_bytes, 2) + + raise ValueError( + f"Free memory on device " + f"({GiB(self.init_snapshot.free_memory)}/" + f"{GiB(self.init_snapshot.total_memory)} GiB) on startup " + f"is less than desired GPU memory utilization " + f"({self.cache_config.gpu_memory_utilization}, " + f"{GiB(self.requested_memory)} GiB). Decrease GPU memory " + f"utilization or reduce GPU memory used by other processes." + ) + else: + raise RuntimeError(f"Not support device type: {self.device_config.device}") + + # Construct the model runner + self.model_runner: GPUGenerationModelRunner = GPUGenerationModelRunner(self.vllm_config, self.device) + + if self.rank == 0: + # If usage stat is enabled, collect relevant info. + report_usage_stats(self.vllm_config) diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index e9de36726..f2d9c1111 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Optional, Union, cast +from typing import TYPE_CHECKING, Any, Optional, Union, cast import numpy as np import torch @@ -6,15 +6,18 @@ from vllm.config import CUDAGraphMode from vllm.distributed.parallel_state import get_pp_group from vllm.forward_context import BatchDescriptor, set_forward_context from vllm.logger import init_logger -from vllm.model_executor.models.interfaces import supports_mrope from vllm.model_executor.models.interfaces_base import VllmModelForPooling from vllm.sampling_params import SamplingType from vllm.utils import LazyLoader, cdiv -from vllm.v1.attention.backends.utils import CommonAttentionMetadata, split_attn_metadata +from vllm.v1.attention.backends.utils import ( + CommonAttentionMetadata, + split_attn_metadata, +) from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.worker.gpu_input_batch import CachedRequestState -from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata +from vllm.v1.worker.gpu_model_runner import GPUModelRunner, IntermediateTensors, PerLayerAttnMetadata from vllm.v1.worker.ubatch_splitting import ubatch_split +from vllm.v1.worker.ubatch_utils import UBatchSlices if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -30,61 +33,6 @@ logger = init_logger(__name__) class OmniGPUModelRunner(GPUModelRunner): - """GPU model runner for omni models with M-RoPE support. - - Extends the base GPUModelRunner to handle multimodal rotary position - embeddings (M-RoPE) for processing image, video, and audio inputs - with proper positional encoding. - """ - - def _init_mrope_positions(self, req_state: CachedRequestState): - """Initialize M-RoPE positions for multimodal inputs. - - Extracts multimodal feature metadata (image grids, video grids, - audio features) and computes M-RoPE positions for proper positional - encoding of multimodal tokens. - - Args: - req_state: Cached request state containing multimodal features - - Raises: - AssertionError: If the model does not support M-RoPE - """ - image_grid_thw = [] - video_grid_thw = [] - second_per_grid_ts = [] - audio_feature_lengths = [] - use_audio_in_video = False - for mm_feature in req_state.mm_features: - mm_item = mm_feature.data - if mm_item is None: - continue - mm_input = mm_item.get_data() - if (t := mm_input.get("image_grid_thw")) is not None: - image_grid_thw.append(t.tolist()) - if (t := mm_input.get("video_grid_thw")) is not None: - video_grid_thw.append(t.tolist()) - if (t := mm_input.get("second_per_grid_ts")) is not None: - second_per_grid_ts.append(t) - if (t := mm_input.get("audio_feature_lengths")) is not None: - audio_feature_lengths.append(t) - # Check for use_audio_in_video - use_audio_in_video_value = mm_input.get("use_audio_in_video") - if use_audio_in_video_value is not None: - use_audio_in_video = bool(use_audio_in_video_value.item()) - - assert supports_mrope(self.get_model()), "M-RoPE support is not implemented." - - req_state.mrope_positions, req_state.mrope_position_delta = self.model.get_mrope_input_positions( - req_state.prompt_token_ids, - hf_config=self.model_config.hf_config, - image_grid_thw=image_grid_thw, - video_grid_thw=video_grid_thw, - second_per_grid_ts=second_per_grid_ts, - audio_feature_lengths=audio_feature_lengths, - use_audio_in_video=use_audio_in_video, - ) - def _update_states(self, scheduler_output: "SchedulerOutput") -> None: """Update the cached states and the persistent batch with the scheduler output. @@ -600,3 +548,291 @@ class OmniGPUModelRunner(GPUModelRunner): logit_indices = np.cumsum(num_scheduled_tokens) - 1 hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) return hidden_states, hidden_states[logit_indices] + + def _preprocess( + self, + scheduler_output: "SchedulerOutput", + num_scheduled_tokens_np: np.ndarray, + intermediate_tensors: Optional[IntermediateTensors] = None, + ubatch_slices: Optional[UBatchSlices] = None, + num_tokens_after_padding: Optional[torch.Tensor] = None, + ) -> tuple[ + int, + int, + Optional[torch.Tensor], + Optional[torch.Tensor], + Optional[torch.Tensor], + torch.Tensor, + Optional[IntermediateTensors], + dict[str, Any], + Optional[dict[str, dict]], + ]: + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + if ubatch_slices: + assert num_tokens_after_padding is not None + num_input_tokens = int(num_tokens_after_padding[0].item() * 2) + self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens) + elif ubatch_slices is None: + num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens) + num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens) + num_input_tokens += num_pad + + # _prepare_inputs may reorder the batch, so we must gather multi + # modal outputs after that to ensure the correct order + per_req_additional_information: Optional[dict[str, dict]] = None + if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder: + # Build multimodal inputs and overlay prompt embeds; collect per-request info + per_req_additional_information = self._build_mm_inputs_and_overlays( + scheduler_output, num_scheduled_tokens, num_scheduled_tokens_np + ) + + input_ids = self.input_ids.gpu[:num_input_tokens] + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = { + **self._init_model_kwargs(num_scheduled_tokens), + **self._extract_mm_kwargs(scheduler_output), + } + elif self.enable_prompt_embeds and get_pp_group().is_first_rank: + # Get the input embeddings for the tokens that are not input embeds, + # then put them into the appropriate positions. + # TODO(qthequartermasterman): Since even when prompt embeds are + # enabled, (a) not all requests will use prompt embeds, and (b) + # after the initial prompt is processed, the rest of the generated + # tokens will be token ids, it is not desirable to have the + # embedding layer outside of the CUDA graph all the time. The v0 + # engine avoids this by "double compiling" the CUDA graph, once + # with input_ids and again with inputs_embeds, for all num_tokens. + # If a batch only has token ids, then including the embedding layer + # in the CUDA graph will be more performant (like in the else case + # below). + token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1) + # Some tokens ids may need to become embeds + if token_ids_idx.numel() > 0: + token_ids = self.input_ids.gpu[token_ids_idx] + tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids) + self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds + + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = self._init_model_kwargs(num_input_tokens) + input_ids = None + else: + # For text-only models, we use token ids as input. + # While it is possible to use embeddings as input just like the + # multimodal models, it is not desirable for performance since + # then the embedding layer is not included in the CUDA graph. + input_ids = self.input_ids.gpu[:num_input_tokens] + inputs_embeds = None + model_kwargs = self._init_model_kwargs(num_input_tokens) + if self.uses_mrope: + positions = self.mrope_positions.gpu[:, :num_input_tokens] + else: + positions = self.positions.gpu[:num_input_tokens] + + if get_pp_group().is_first_rank: + intermediate_tensors = None + else: + intermediate_tensors = self.sync_and_slice_intermediate_tensors( + num_input_tokens, intermediate_tensors, True + ) + + if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: + encoder_inputs = self._extract_encoder_inputs(scheduler_output) + model_kwargs.update(encoder_inputs) + + return ( + num_scheduled_tokens, + num_input_tokens, + num_tokens_after_padding, + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + per_req_additional_information, + ) + + def _decode_and_store_request_payloads(self, scheduler_output: "SchedulerOutput") -> None: + """Decode per-request prompt_embeds and additional_information for newly + scheduled requests and store them to CPU in the request state. + This version avoids hard dependency on payload classes by duck-typing.""" + try: + new_reqs = getattr(scheduler_output, "scheduled_new_reqs", []) + if not new_reqs: + return + for nr in new_reqs: + req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None) + if req_id is None: + continue + # prompt_embeds + payload_pe = getattr(nr, "prompt_embeds", None) + pe_cpu = None + if payload_pe is not None: + if isinstance(payload_pe, torch.Tensor): + pe_cpu = payload_pe.detach().to("cpu").contiguous() + else: + # Try duck-typing a payload with data/shape/dtype + data = getattr(payload_pe, "data", None) + shape = getattr(payload_pe, "shape", None) + if data is not None and shape is not None: + dt = np.dtype(getattr(payload_pe, "dtype", "float32")) + arr = np.frombuffer(data, dtype=dt) + arr = arr.reshape(shape) + pe_cpu = torch.from_numpy(arr) + if pe_cpu is not None and req_id in self.requests: + setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu) + # additional_information + payload_info = getattr(nr, "additional_information", None) + if payload_info is not None: + info_dict = {} + if isinstance(payload_info, dict): + info_dict = payload_info + else: + # Try duck-typing a payload with entries, each entry may have + # tensor_data/tensor_dtype/tensor_shape or list_data + entries = getattr(payload_info, "entries", None) + if isinstance(entries, dict): + for k, entry in entries.items(): + tensor_data = getattr(entry, "tensor_data", None) + if tensor_data is not None: + dt = np.dtype(getattr(entry, "tensor_dtype", "float32")) + arr = np.frombuffer(tensor_data, dtype=dt) + arr = arr.reshape(getattr(entry, "tensor_shape", ())) + info_dict[k] = torch.from_numpy(arr) + else: + info_dict[k] = getattr(entry, "list_data", None) + if info_dict and req_id in self.requests: + setattr(self.requests[req_id], "additional_information_cpu", info_dict) + except Exception as e: + logger.error(f"Error decoding prompt_embeds / additional_information: {e}") + + def _gather_runtime_additional_information(self) -> list[dict]: + """Gather per-request additional_information stored in request state in batch order.""" + per_req_runtime_info = [] + for req_id in self.input_batch.req_ids: + req_state = self.requests.get(req_id) + info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None + per_req_runtime_info.append(info if isinstance(info, dict) else {}) + return per_req_runtime_info + + def _compute_request_token_spans(self, num_scheduled_tokens_np) -> list[tuple[int, int]]: + """Compute (start, end) token spans for each request within the flattened step sequence.""" + req_token_spans: list[tuple[int, int]] = [] + for req_index in range(len(self.input_batch.req_ids)): + start_offset = int(self.query_start_loc.cpu[req_index]) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + req_token_spans.append((start_offset, start_offset + sched_tokens)) + return req_token_spans + + def _build_model_kwargs_extra( + self, + per_req_additional_information: Optional[dict[str, dict]], + num_scheduled_tokens_np, + ) -> dict: + """Build extra keyword arguments passed to the model for this step, including: + - additional_information_by_req_id: per-request additional information provided by preprocess for this step + - runtime_additional_information: per-request additional information stored in request state + - request_ids: the request id list in batch order + - request_token_spans: token spans per request within the flattened sequence + """ + model_kwargs_extra: dict[str, object] = {} + if per_req_additional_information: + model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information + try: + model_kwargs_extra["runtime_additional_information"] = self._gather_runtime_additional_information() + model_kwargs_extra["request_ids"] = self.input_batch.req_ids + model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(num_scheduled_tokens_np) + except Exception: + pass + return model_kwargs_extra + + def _process_additional_information_updates(self, multimodal_outputs: object) -> None: + """Process model-provided per-request additional_information updates and merge into request state.""" + try: + if not isinstance(multimodal_outputs, dict): + return + if ( + "additional_information_update" not in multimodal_outputs + and "additional_information_update_by_req_id" not in multimodal_outputs + ): + return + updates_list = multimodal_outputs.get("additional_information_update") + if isinstance(updates_list, list): + for idx, upd in enumerate(updates_list): + if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids): + continue + req_id = self.input_batch.req_ids[idx] + self._merge_additional_information_update(req_id, upd) + updates_map = multimodal_outputs.get("additional_information_update_by_req_id") + if isinstance(updates_map, dict): + for req_id, upd in updates_map.items(): + if not isinstance(upd, dict): + continue + if req_id not in self.requests: + continue + self._merge_additional_information_update(req_id, upd) + 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}" + ) + + def _build_mm_inputs_and_overlays( + self, + scheduler_output: "SchedulerOutput", + num_scheduled_tokens: int, + num_scheduled_tokens_np: np.ndarray, + ) -> dict[str, dict]: + """Run MM encoder, fill scheduled inputs_embeds, and overlay prompt embeds; + collect per-request additional_information for the prefill portion.""" + # Run the multimodal encoder if any. + self._execute_mm_encoder(scheduler_output) + mm_embeds = self._gather_mm_embeddings(scheduler_output) + # Always use embeddings as inputs for MM models + inputs_embeds_scheduled = self.model.get_input_embeddings( + input_ids=self.input_ids.gpu[:num_scheduled_tokens], + multimodal_embeddings=mm_embeds or None, + ) + self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) + # Reset per-step additional information collector (deprecated concat path) + if hasattr(self, "_forward_additional_information"): + self._forward_additional_information = None + # Overlay and collect per-request info + return self._collect_additional_information_for_prefill(num_scheduled_tokens_np) + + def _collect_additional_information_for_prefill( + self, + num_scheduled_tokens_np: np.ndarray, + ) -> dict[str, dict]: + """Overlay per-request prompt_embeds for the prefill portion and collect + additional_information slices for this step. Returns a map req_id -> dict.""" + per_req_additional_information: dict[str, dict] = {} + for req_index, req_id in enumerate(self.input_batch.req_ids): + req_state = self.requests[req_id] + pe_cpu = getattr(req_state, "prompt_embeds_cpu", None) + addi_cpu = getattr(req_state, "additional_information_cpu", None) + num_computed_tokens = int(self.input_batch.num_computed_tokens_cpu[req_index]) + prompt_len = len(req_state.prompt_token_ids) + prompt_remaining = max(0, prompt_len - num_computed_tokens) + sched_tokens = int(num_scheduled_tokens_np[req_index]) + overlay_len = min(sched_tokens, prompt_remaining) + if overlay_len <= 0 and not (isinstance(addi_cpu, dict) and addi_cpu): + continue + if overlay_len > 0 and pe_cpu is not None: + src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to( + dtype=self.dtype, device=self.device, non_blocking=True + ) + start_offset = int(self.query_start_loc.cpu[req_index]) + self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src) + if addi_cpu is not None and isinstance(addi_cpu, dict): + req_info: dict[str, object] = {} + for k, v in addi_cpu.items(): + if isinstance(v, torch.Tensor): + req_info[k] = v.detach().to("cpu").contiguous() + elif isinstance(v, list): + req_info[k] = v + else: + req_info[k] = v + if req_info: + per_req_additional_information[req_id] = req_info + return per_req_additional_information