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

Merge branch 'main' into feat/hidden-state-input-request-input-processor-patch

Browse Source 
Signed-off-by: Zhou Taichang <tzhouam@connect.ust.hk>
This commit is contained in:
Zhou Taichang authored and GitHub committed 2025-10-24 11:52:45 +08:00
2 parents 63c72f4c2b 4a4c3c1b82
commit 582639c77a
31 files changed
+8400 -1100
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

README.md
+27 -98
View File
@@ -8,7 +8,7 @@ Traditional vLLM systems are limited to text-based, autoregressive generation. v
- **Multi-modal Models**: Text, image, video, audio, and sensor data processing
- **Non-autoregressive Architectures**: Diffusion Transformers (DiT) and other parallel generation models
- **Heterogeneous Outputs**: Beyond traditional text generation to structured, binary, and streaming outputs
- **Heterogeneous Outputs**: Beyond traditional text generation to multimodal outputs
## 🏗️ Architecture
@@ -28,119 +28,48 @@ vLLM-omni is built on a modular architecture that extends vLLM's core functional
- **Text**: Advanced tokenization and embedding generation
- **Image**: Vision encoder integration (CLIP, etc.)
- **Audio**: Speech processing and audio embedding
- **Video**: Frame-by-frame and temporal processing
- **Sensor**: IoT and sensor data interpretation
### Output Formats
- **Structured Data**: JSON, XML, and custom formats
- **Binary Outputs**: Images, audio, and video generation
- **Streaming**: Real-time progressive generation
- **Multipart**: Combined multi-modal responses
## 📋 Supported Models
### AR + Diffusion Transformer (DiT) Models
- Qwen-Image (Image generation and editing)
- Qwen-omni (Thinker-Talker-Codec structure)
- Custom DiT and hiybrid architectures
- HunyunaImage 3.0 (Ongoing)
- Qwen-Image (Ongoing)
## 🛠️ Installation
### Quick Start
#### Option 1: Docker (Recommended for macOS)
Set up basic environments
```bash
uv venv --python 3.12 --seed
source .venv/bin/activate
```
Install certain version of vllm with commitid: 808a7b69df479b6b3a16181711cac7ca28a9b941
```bash
# Clone the repository
git clone https://github.com/hsliuustc0106/vllm-omni.git
cd vllm-omni
# Run the automated Docker setup
./scripts/docker-setup-macos.sh
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout 808a7b69df479b6b3a16181711cac7ca28a9b941
VLLM_USE_PRECOMPILED=1 uv pip install --editable .
```
#### Option 2: Local Installation
## Run examples (Qwen2.5-omni)
Get into the example folder
```bash
# Clone the repository
git clone https://github.com/hsliuustc0106/vllm-omni.git
cd vllm-omni
# Run the installation script
./install.sh
cd vllm_omni
cd examples/offline_inference/qwen2_5_omni
```
### Prerequisites
- Python 3.11+ (recommended)
- Conda or Miniconda
- Git
- CUDA 11.8+ (for GPU acceleration) or CPU-only installation
### Installation Methods
#### Method 1: Automated Installation (Recommended)
Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run.
```bash
# Using shell script
./install.sh
# Or using Python script
python install.py
bash run.sh
```
The output audio is saved in ./output_audio
#### Method 2: Manual Installation
```bash
# Create conda environment
conda create -n vllm_omni python=3.11 -y
conda activate vllm_omni
## To-do list
- [x] Offline inference example for Qwen2.5-omni with single request
- [ ] Adaptation from current vllm branch to stable vllm v0.11.0
- [ ] Offline inference example for Qwen2.5-omni with streaming multiple requests
- [ ] Online inference support
- [ ] Support for other models
# Install PyTorch (CPU or GPU)
pip install torch>=2.7 --index-url https://download.pytorch.org/whl/cpu # CPU
# pip install torch>=2.7 --index-url https://download.pytorch.org/whl/cu121 # GPU
# Install dependencies
pip install -r requirements.txt
pip install "vllm>=0.10.2"
# Install vLLM-omni
pip install -e .
```
### Verify Installation
```bash
# Test the installation
python test_installation.py
# Test basic functionality
python -c "import vllm_omni; print('Ready!')"
# Test CLI
vllm --help
```
For detailed installation instructions, see [INSTALL.md](INSTALL.md).
## 📥 Model Download
Models are automatically downloaded when first used, or you can pre-download them:
```bash
# Check downloaded models
python scripts/download_models.py --check-cache
# Download all default models
python scripts/download_models.py --all
# Download specific models
python scripts/download_models.py --ar-models Qwen/Qwen3-0.6B
python scripts/download_models.py --dit-models stabilityai/stable-diffusion-2-1
```
**Model Storage Location:**
- Default: `~/.cache/huggingface/hub/`
- AR models: 100MB - 1GB each
- DiT models: 2GB - 7GB each
For detailed model management, see [MODEL_DOWNLOAD_GUIDE.md](docs/MODEL_DOWNLOAD_GUIDE.md).
For detailed model management, see [vllm_omni_design.md](docs/architecture/vllm_omni_design.md) and [high_level_arch_design.md](docs/architecture/high_level_arch_design.md).
examples/offline_inference/qwen_2_5_omni/README.md
+37
View File
@@ -0,0 +1,37 @@
# Offline Example of vLLM-omni for Qwen2.5-omni
## Installation
Set up basic environments
```bash
uv venv --python 3.12 --seed
source .venv/bin/activate
```
Install certain version of vllm with commitid: 808a7b69df479b6b3a16181711cac7ca28a9b941
```bash
git clone https://github.com/vllm-project/vllm.git
cd vllm
git checkout 808a7b69df479b6b3a16181711cac7ca28a9b941
VLLM_USE_PRECOMPILED=1 uv pip install --editable .
```
## Run examples
Get into the example folder
```bash
cd vllm_omni
cd examples/offline_inference/qwen2_5_omni
```
Modify PYTHONPATH in run.sh as your path of vllm_omni. Then run.
```bash
bash run.sh
```
The output audio is saved in ./output_audio
## To-do list
- [x] Offline inference example for Qwen2.5-omni with single request
- [ ] Adaptation from current vllm branch to stable vllm v0.11.0
- [ ] Offline inference example for Qwen2.5-omni with streaming multiple requests
- [ ] Online inference support
- [ ] Support for other models
examples/offline_inference/qwen_2_5_omni/end2end.py
+130
View File
@@ -0,0 +1,130 @@
import argparse
import os
import soundfile as sf
import random
import numpy as np
import torch
from vllm.sampling_params import SamplingParams
import os as _os_env_toggle
_os_env_toggle.environ["VLLM_USE_V1"] = "1"
from vllm_omni.entrypoints.omni_llm import OmniLLM
from utils import make_omni_prompt
SEED = 42
# Set all random seeds
random.seed(SEED)
np.random.seed(SEED)
torch.manual_seed(SEED)
torch.cuda.manual_seed(SEED)
torch.cuda.manual_seed_all(SEED)
# Make PyTorch deterministic
torch.backends.cudnn.deterministic = True
torch.backends.cudnn.benchmark = False
# Set environment variables for deterministic behavior
os.environ["PYTHONHASHSEED"] = str(SEED)
os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
def parse_args():
parser = argparse.ArgumentParser()
parser.add_argument('--model', required=True, help='Path to merged model directory (will be created if downloading).')
parser.add_argument('--thinker-model', type=str, default=None)
parser.add_argument('--talker-model', type=str, default=None)
parser.add_argument('--code2wav-model', type=str, default=None)
parser.add_argument('--hf-hub-id', default='Qwen/Qwen2.5-Omni-7B', help='Hugging Face repo id to download if needed.')
parser.add_argument('--hf-revision', default=None, help='Optional HF revision (branch/tag/commit).')
parser.add_argument('--prompts', required=True, nargs='+', help='Input text prompts.')
parser.add_argument('--voice-type', default='default', help='Voice type, e.g., m02, f030, default.')
parser.add_argument('--code2wav-dir', default=None, help='Path to code2wav folder (contains spk_dict.pt).')
parser.add_argument('--dit-ckpt', default=None, help='Path to DiT checkpoint file (e.g., dit.pt).')
parser.add_argument('--bigvgan-ckpt', default=None, help='Path to BigVGAN checkpoint file.')
parser.add_argument('--dtype', default='bfloat16', choices=['float16', 'bfloat16', 'float32'])
parser.add_argument('--max-model-len', type=int, default=32768)
parser.add_argument("--thinker-only", action="store_true")
parser.add_argument("--text-only", action="store_true")
parser.add_argument("--do-wave", action="store_true")
parser.add_argument('--prompt_type',
choices=[
'text', 'audio', 'audio-long', 'audio-long-chunks',
'audio-long-expand-chunks', 'image', 'video',
'video-frames', 'audio-in-video', 'audio-in-video-v2',
"audio-multi-round", "badcase-vl", "badcase-text",
"badcase-image-early-stop", "badcase-two-audios",
"badcase-two-videos", "badcase-multi-round",
"badcase-voice-type", "badcase-voice-type-v2",
"badcase-audio-tower-1", "badcase-audio-only"
],
default='text')
parser.add_argument('--use-torchvision', action='store_true')
parser.add_argument('--tokenize', action='store_true')
parser.add_argument('--output-wav', default="output.wav", help='Output wav file path.')
parser.add_argument('--thinker-hidden-states-dir', default="thinker_hidden_states", help='Path to thinker hidden states directory.')
args = parser.parse_args()
return args
def main():
args = parse_args()
model_name = args.model
omni_llm = OmniLLM(model=model_name)
thinker_sampling_params = SamplingParams(
temperature=0.0, # Deterministic - no randomness
top_p=1.0, # Disable nucleus sampling
top_k=-1, # Disable top-k sampling
max_tokens=2048,
seed=SEED, # Fixed seed for sampling
detokenize=True,
repetition_penalty=1.1,
)
talker_sampling_params = SamplingParams(
temperature=0.0, # Deterministic - no randomness
top_p=1.0, # Disable nucleus sampling
top_k=-1, # Disable top-k sampling
max_tokens=2048,
seed=SEED, # Fixed seed for sampling
detokenize=True,
repetition_penalty=1.1,
stop_token_ids=[8294]
)
code2wav_sampling_params = SamplingParams(
temperature=0.0, # Deterministic - no randomness
top_p=1.0, # Disable nucleus sampling
top_k=-1, # Disable top-k sampling
max_tokens=2048,
seed=SEED, # Fixed seed for sampling
detokenize=True,
repetition_penalty=1.1,
)
sampling_params_list = [thinker_sampling_params,
talker_sampling_params,
code2wav_sampling_params]
prompt = [make_omni_prompt(args, prompt) for prompt in args.prompts]
omni_outputs = omni_llm.generate(prompt, sampling_params_list)
os.makedirs(args.output_wav, 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 = output.request_id
text_output = output.outputs[0].text
print(f"Request ID: {request_id}, Text Output: {text_output}")
elif stage_outputs.final_output_type == "audio":
for output in stage_outputs.request_output:
request_id = output.request_id
audio_tensor = output.multimodal_output["audio"]
output_wav = os.path.join(args.output_wav, f"output_{output.request_id}.wav")
sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
print(f"Request ID: {request_id}, Saved audio to {output_wav}")
if __name__ == "__main__":
main()
examples/offline_inference/qwen_2_5_omni/output_audio/output_0.wav
BIN
View File
Binary file not shown.
examples/offline_inference/qwen_2_5_omni/processing_omni.py
+373
View File
@@ -0,0 +1,373 @@
from __future__ import annotations
import base64
import logging
import math
import os
import time
import warnings
from functools import lru_cache
from io import BytesIO
import requests
import torch
import torchvision
from packaging import version
from PIL import Image
from torchvision import io, transforms
from torchvision.transforms import InterpolationMode
logger = logging.getLogger(__name__)
IMAGE_FACTOR = 28
MIN_PIXELS = 4 * 28 * 28
MAX_PIXELS = 16384 * 28 * 28
MAX_RATIO = 200
VIDEO_MIN_PIXELS = 128 * 28 * 28
VIDEO_MAX_PIXELS = 768 * 28 * 28
VIDEO_TOTAL_PIXELS = 24576 * 28 * 28
FRAME_FACTOR = 2
FPS = 2.0
FPS_MIN_FRAMES = 4
FPS_MAX_FRAMES = 768
temporal_patch_size = 2
spatial_patch_size = 14
spatial_merge_size = 2
def round_by_factor(number: int, factor: int) -> int:
"""Returns the closest integer to 'number' that is divisible by 'factor'."""
return round(number / factor) * factor
def ceil_by_factor(number: int, factor: int) -> int:
"""Returns the smallest integer greater than or equal to 'number' that is divisible by 'factor'."""
return math.ceil(number / factor) * factor
def floor_by_factor(number: int, factor: int) -> int:
"""Returns the largest integer less than or equal to 'number' that is divisible by 'factor'."""
return math.floor(number / factor) * factor
def smart_resize(height: int,
width: int,
factor: int = IMAGE_FACTOR,
min_pixels: int = MIN_PIXELS,
max_pixels: int = MAX_PIXELS) -> tuple[int, int]:
"""
Rescales the image so that the following conditions are met:
1. Both dimensions (height and width) are divisible by 'factor'.
2. The total number of pixels is within the range ['min_pixels', 'max_pixels'].
3. The aspect ratio of the image is maintained as closely as possible.
"""
if max(height, width) / min(height, width) > MAX_RATIO:
raise ValueError(
f"absolute aspect ratio must be smaller than {MAX_RATIO}, got {max(height, width) / min(height, width)}"
)
h_bar = max(factor, round_by_factor(height, factor))
w_bar = max(factor, round_by_factor(width, factor))
if h_bar * w_bar > max_pixels:
beta = math.sqrt((height * width) / max_pixels)
h_bar = floor_by_factor(height / beta, factor)
w_bar = floor_by_factor(width / beta, factor)
elif h_bar * w_bar < min_pixels:
beta = math.sqrt(min_pixels / (height * width))
h_bar = ceil_by_factor(height * beta, factor)
w_bar = ceil_by_factor(width * beta, factor)
return h_bar, w_bar
def fetch_image(ele: dict[str, str | Image.Image],
size_factor: int = IMAGE_FACTOR) -> Image.Image:
if "image" in ele:
image = ele["image"]
else:
image = ele["image_url"]
image_obj = None
if isinstance(image, Image.Image):
image_obj = image
elif image.startswith("http://") or image.startswith("https://"):
image_obj = Image.open(requests.get(image, stream=True).raw)
elif image.startswith("file://"):
image_obj = Image.open(image[7:])
elif image.startswith("data:image"):
if "base64," in image:
_, base64_data = image.split("base64,", 1)
data = base64.b64decode(base64_data)
image_obj = Image.open(BytesIO(data))
else:
image_obj = Image.open(image)
if image_obj is None:
raise ValueError(
f"Unrecognized image input, support local path, http url, base64 and PIL.Image, got {image}"
)
image = image_obj.convert("RGB")
## resize
if "resized_height" in ele and "resized_width" in ele:
resized_height, resized_width = smart_resize(
ele["resized_height"],
ele["resized_width"],
factor=size_factor,
)
else:
width, height = image.size
min_pixels = ele.get("min_pixels", MIN_PIXELS)
max_pixels = ele.get("max_pixels", MAX_PIXELS)
resized_height, resized_width = smart_resize(
height,
width,
factor=size_factor,
min_pixels=min_pixels,
max_pixels=max_pixels,
)
image = image.resize((resized_width, resized_height))
return image
def smart_nframes(
ele: dict,
total_frames: int,
video_fps: int | float,
) -> int:
"""calculate the number of frames for video used for model inputs.
Args:
ele (dict): a dict contains the configuration of video.
support either `fps` or `nframes`:
- nframes: the number of frames to extract for model inputs.
- fps: the fps to extract frames for model inputs.
- min_frames: the minimum number of frames of the video, only used when fps is provided.
- max_frames: the maximum number of frames of the video, only used when fps is provided.
total_frames (int): the original total number of frames of the video.
video_fps (int | float): the original fps of the video.
Raises:
ValueError: nframes should in interval [FRAME_FACTOR, total_frames].
Returns:
int: the number of frames for video used for model inputs.
"""
assert not ("fps" in ele
and "nframes" in ele), "Only accept either `fps` or `nframes`"
if "nframes" in ele:
nframes = round_by_factor(ele["nframes"], FRAME_FACTOR)
else:
fps = ele.get("fps", FPS)
min_frames = ceil_by_factor(ele.get("min_frames", FPS_MIN_FRAMES),
FRAME_FACTOR)
max_frames = floor_by_factor(
ele.get("max_frames", min(FPS_MAX_FRAMES, total_frames)),
FRAME_FACTOR)
nframes = total_frames / video_fps * fps
nframes = min(max(nframes, min_frames), max_frames)
nframes = round_by_factor(nframes, FRAME_FACTOR)
if not (FRAME_FACTOR <= nframes and nframes <= total_frames):
raise ValueError(
f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}."
)
return nframes
def _read_video_torchvision(ele: dict, ) -> torch.Tensor:
"""read video using torchvision.io.read_video
Args:
ele (dict): a dict contains the configuration of video.
support keys:
- video: the path of video. support "file://", "http://", "https://" and local path.
- video_start: the start time of video.
- video_end: the end time of video.
Returns:
torch.Tensor: the video tensor with shape (T, C, H, W).
"""
video_path = ele["video"]
if version.parse(torchvision.__version__) < version.parse("0.19.0"):
if "http://" in video_path or "https://" in video_path:
warnings.warn(
"torchvision < 0.19.0 does not support http/https video path, please upgrade to 0.19.0."
)
if "file://" in video_path:
video_path = video_path[7:]
st = time.time()
video, audio, info = io.read_video(
video_path,
start_pts=ele.get("video_start", 0.0),
end_pts=ele.get("video_end", None),
pts_unit="sec",
output_format="TCHW",
)
total_frames, video_fps = video.size(0), info["video_fps"]
total_duration = round(total_frames / video_fps, 3)
logger.info(
f"torchvision: {video_path=}, {total_frames=}, {video_fps=}, duration={total_duration}s, time={time.time() - st:.3f}s"
)
nframes = smart_nframes(ele,
total_frames=total_frames,
video_fps=video_fps)
idx = torch.linspace(0, total_frames - 1, nframes).round().long()
video = video[idx]
return video, total_duration, nframes
def is_decord_available() -> bool:
import importlib.util
return importlib.util.find_spec("decord") is not None
def _read_video_decord(ele: dict, ) -> torch.Tensor:
"""read video using decord.VideoReader
Args:
ele (dict): a dict contains the configuration of video.
support keys:
- video: the path of video. support "file://", "http://", "https://" and local path.
- video_start: the start time of video.
- video_end: the end time of video.
Returns:
torch.Tensor: the video tensor with shape (T, C, H, W).
"""
import decord
video_path = ele["video"]
st = time.time()
vr = decord.VideoReader(video_path)
# TODO: support start_pts and end_pts
if 'video_start' in ele or 'video_end' in ele:
raise NotImplementedError(
"not support start_pts and end_pts in decord for now.")
total_frames, video_fps = len(vr), vr.get_avg_fps()
total_duration = round(total_frames / video_fps, 3)
logger.info(
f"decord: {video_path=}, {total_frames=}, {video_fps=}, time={time.time() - st:.3f}s"
)
nframes = smart_nframes(ele,
total_frames=total_frames,
video_fps=video_fps)
idx = torch.linspace(0, total_frames - 1, nframes).round().long().tolist()
video = vr.get_batch(idx).asnumpy()
video = torch.tensor(video).permute(0, 3, 1, 2) # Convert to TCHW format
return video, total_duration, nframes
VIDEO_READER_BACKENDS = {
"decord": _read_video_decord,
"torchvision": _read_video_torchvision,
}
FORCE_QWENVL_VIDEO_READER = os.getenv("FORCE_QWENVL_VIDEO_READER", None)
@lru_cache(maxsize=1)
def get_video_reader_backend() -> str:
if FORCE_QWENVL_VIDEO_READER is not None:
video_reader_backend = FORCE_QWENVL_VIDEO_READER
elif is_decord_available():
video_reader_backend = "decord"
else:
video_reader_backend = "torchvision"
# print(f"qwen-vl-utils using {video_reader_backend} to read video.", file=sys.stderr)
return video_reader_backend
def fetch_video(
ele: dict,
image_factor: int = IMAGE_FACTOR) -> torch.Tensor | list[Image.Image]:
if isinstance(ele["video"], str):
video_reader_backend = get_video_reader_backend()
video, total_dur, nframes = VIDEO_READER_BACKENDS[
video_reader_backend](ele)
frame_timestamps = total_dur * torch.arange(1, nframes + 1) / nframes
grid_timestamps = frame_timestamps[::FRAME_FACTOR]
second_per_grid = grid_timestamps[1] - grid_timestamps[0]
nframes, _, height, width = video.shape
factor = spatial_patch_size * spatial_merge_size
min_pixels = ele.get("min_pixels", VIDEO_MIN_PIXELS)
total_pixels = ele.get("total_pixels", VIDEO_TOTAL_PIXELS)
max_pixels = max(
min(VIDEO_MAX_PIXELS, total_pixels / nframes * FRAME_FACTOR),
int(min_pixels * 1.05))
max_pixels = ele.get("max_pixels", max_pixels)
# min_pixels = (factor ** 2) * 52
# max_pixels = (factor ** 2) * min(768, (16384 / nframes * temporal_patch_size))
if "resized_height" in ele and "resized_width" in ele:
resized_height, resized_width = smart_resize(
ele["resized_height"],
ele["resized_width"],
factor=image_factor,
)
else:
resized_height, resized_width = smart_resize(
height,
width,
factor=image_factor,
min_pixels=min_pixels,
max_pixels=max_pixels,
)
video = transforms.functional.resize(
video,
[resized_height, resized_width],
interpolation=InterpolationMode.BICUBIC,
antialias=True,
).float()
return video, total_dur, nframes, second_per_grid
else:
assert isinstance(ele["video"], (list, tuple))
process_info = ele.copy()
process_info.pop("type", None)
process_info.pop("video", None)
images = [
fetch_image({
"image": video_element,
**process_info
},
size_factor=image_factor)
for video_element in ele["video"]
]
nframes = ceil_by_factor(len(images), FRAME_FACTOR)
if len(images) < nframes:
images.extend([images[-1]] * (nframes - len(images)))
return images, None, None, None
def extract_vision_info(
conversations: list[dict] | list[list[dict]]) -> list[dict]:
vision_infos = []
if isinstance(conversations[0], dict):
conversations = [conversations]
for conversation in conversations:
for message in conversation:
if isinstance(message["content"], list):
for ele in message["content"]:
if ("image" in ele or "image_url" in ele or "video" in ele
or ele["type"] in ("image", "image_url", "video")):
vision_infos.append(ele)
return vision_infos
def process_vision_info(
conversations: list[dict] | list[list[dict]],
) -> tuple[list[Image.Image] | None, list[torch.Tensor | list[Image.Image]]
| None]:
vision_infos = extract_vision_info(conversations)
## Read images or videos
image_inputs = []
video_inputs = []
for vision_info in vision_infos:
if "image" in vision_info or "image_url" in vision_info:
image_inputs.append(fetch_image(vision_info))
elif "video" in vision_info:
video_inputs.append(fetch_video(vision_info))
else:
raise ValueError("image, image_url or video should in content.")
if len(image_inputs) == 0:
image_inputs = None
if len(video_inputs) == 0:
video_inputs = None
return image_inputs, video_inputs
examples/offline_inference/qwen_2_5_omni/run.sh
+9
View File
@@ -0,0 +1,9 @@
export PYTHONPATH=/path/to/vllm-omni:$PYTHONPATH
export HF_ENDPOINT=https://hf-mirror.com
python end2end.py --model Qwen/Qwen2.5-Omni-7B \
--prompts "Explain the system architecture for a scalable audio generation pipeline. Answer in 15 words." \
--voice-type "m02" \
--dit-ckpt none \
--bigvgan-ckpt none \
--output-wav output_audio \
--prompt_type text
examples/offline_inference/qwen_2_5_omni/utils.py
+305
View File
@@ -0,0 +1,305 @@
import tempfile
from urllib.request import urlopen
import librosa
import soundfile as sf
import resampy
from typing import Dict, Optional
import torch
import requests
import torchvision.io
from typing import Union, List
from vllm.inputs import TextPrompt
from vllm_omni.inputs.data import OmniTokensPrompt
from processing_omni import fetch_image, fetch_video
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 resample_wav_to_16khz(input_filepath):
data, original_sample_rate = sf.read(input_filepath)
# Only use the first channel
if len(data.shape) > 1:
data = data[:, 0]
# resample to 16kHz
data_resampled = resampy.resample(data,
sr_orig=original_sample_rate,
sr_new=16000)
return data_resampled
def fetch_and_read_video(args, video_url: str, fps=2):
def read_video_with_torchvision(video_file_name: str):
video, audio, info = torchvision.io.read_video(
video_file_name,
start_pts=0.0,
end_pts=None,
pts_unit="sec",
output_format="TCHW",
)
total_frames, video_fps = video.size(0), info["video_fps"]
total_duration = round(total_frames / video_fps, 3)
nframes = int(total_frames / video_fps * fps)
frame_timestamps = total_duration * torch.arange(1,
nframes + 1) / nframes
grid_timestamps = frame_timestamps[::2]
second_per_grid = grid_timestamps[1] - grid_timestamps[0]
idx = torch.linspace(0, video.size(0) - 1, nframes).round().long()
video = video[idx]
if args.legacy_omni_video:
return [video, total_duration, nframes, second_per_grid.item()]
else:
return video
def read_video_with_transformers(video_file_name: Union[str, List[str]]):
video, total_duration, nframes, second_per_grid = fetch_video(
{'video': video_file_name})
if total_duration is None and nframes is None:
nframes = len(video)
total_duration = 0.5 * nframes
second_per_grid = 1.0
if args.legacy_omni_video:
return [video, total_duration, nframes, second_per_grid]
else:
return video
def read_video(video_file_name: str):
if args.use_torchvision:
return read_video_with_torchvision(video_file_name)
else:
return read_video_with_transformers(video_file_name)
if isinstance(video_url, str) and video_url.startswith("http"):
with tempfile.NamedTemporaryFile(delete=True) as temp_video_file:
resp = requests.get(video_url)
assert resp.status_code == requests.codes.ok, f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
temp_video_file.write(urlopen(video_url).read())
temp_video_file_path = temp_video_file.name
video_file_name = temp_video_file_path
return read_video(video_file_name)
else:
video_file_name = video_url
return read_video(video_file_name)
def make_inputs_qwen2_omni(
args,
messages: List[Dict[str, Union[str, List[Dict[str, str]]]]],
use_audio_in_video: Optional[bool] = False,
tokenize: bool = False,
) -> Union[OmniTokensPrompt, TextPrompt]:
from transformers import AutoConfig, AutoProcessor, AutoTokenizer
processor = AutoProcessor.from_pretrained(args.model)
tokenizer = AutoTokenizer.from_pretrained(args.model)
try:
config = AutoConfig.from_pretrained(args.model)
if 'Qwen2_5OmniModel' in config.architectures:
args.legacy_omni_video = False
else:
args.legacy_omni_video = True
except:
args.legacy_omni_video = True
audios, images, videos = [], [], []
for message in messages:
if not isinstance(message['content'], list):
message['content'] = [{
'type': 'text',
'text': message['content'],
}]
index, num_contents = 0, len(message['content'])
while index < num_contents:
ele = message['content'][index]
if 'type' not in ele:
if 'text' in ele:
ele['type'] = 'text'
elif 'audio' in ele:
ele['type'] = 'audio'
elif 'audio_url' in ele:
ele['type'] = 'audio_url'
elif 'image' in ele:
ele['type'] = 'image'
elif 'image_url' in ele:
ele['type'] = 'image_url'
elif 'video' in ele:
ele['type'] = 'video'
elif 'video_url' in ele:
ele['type'] = 'video_url'
else:
raise ValueError(f'Unknown ele: {ele}')
if ele['type'] == 'audio' or ele['type'] == 'audio_url':
if 'audio_url' in ele:
audio_key = 'audio_url'
with tempfile.NamedTemporaryFile(
delete=True) as temp_audio_file:
temp_audio_file.write(urlopen(ele[audio_key]).read())
temp_audio_file_path = temp_audio_file.name
audios.append(
resample_wav_to_16khz(temp_audio_file_path))
ele['audio'] = temp_audio_file_path
elif 'audio' in ele:
audio_key = 'audio'
audios.append(resample_wav_to_16khz(ele[audio_key]))
else:
raise ValueError(f'Unknown ele {ele}')
elif use_audio_in_video and (ele['type'] == 'video'
or ele['type'] == 'video_url'):
# use video as audio as well
if 'video_url' in ele:
audio_key = 'video_url'
with tempfile.NamedTemporaryFile(
delete=True) as temp_video_file:
temp_video_file.write(urlopen(ele[audio_key]).read())
temp_video_file_path = temp_video_file.name
ele[audio_key] = temp_video_file_path
audios.append(
librosa.load(temp_video_file_path, sr=16000)[0])
videos.append(
fetch_and_read_video(args, temp_video_file_path))
ele['video'] = temp_video_file_path
elif 'video' in ele:
audio_key = 'video'
audios.append(librosa.load(ele[audio_key], sr=16000)[0])
videos.append(fetch_and_read_video(args, audio_key))
else:
raise ValueError("Unknown ele {}".format(ele))
# insert a audio after the video
message['content'].insert(index + 1, {
"type": "audio",
"audio": ele[audio_key],
})
# no need to load the added audio again
index += 1
elif ele['type'] == 'video' or ele['type'] == 'video_url':
if 'video_url' in ele:
video_key = 'video_url'
with tempfile.NamedTemporaryFile(
delete=True) as temp_video_file:
temp_video_file.write(urlopen(ele['video_url']).read())
temp_video_file_path = temp_video_file.name
videos.append(fetch_and_read_video(args, temp_video_file))
ele['video'] = temp_video_file_path
else:
video_key = 'video'
videos.append(fetch_and_read_video(args, ele[video_key]))
elif ele['type'] == 'image' or ele['type'] == 'image_url':
images.append(fetch_image(ele))
# move to the next content
index += 1
prompt = processor.apply_chat_template(
messages,
tokenize=tokenize,
add_generation_prompt=True,
add_vision_id=True,
)
audios = audios if len(audios) > 0 else None
images = images if len(images) > 0 else None
videos = videos if len(videos) > 0 else None
multi_modal_data = {}
if audios:
multi_modal_data["audio"] = audios
if images:
multi_modal_data["image"] = images
if videos:
multi_modal_data["video"] = videos
if isinstance(prompt, list) and isinstance(prompt[0], (list, str)):
prompt = prompt[0]
if tokenize:
return OmniTokensPrompt(
prompt_token_ids=prompt,
multi_modal_data=multi_modal_data,
)
else:
return TextPrompt(
prompt=prompt,
multi_modal_data=multi_modal_data,
)
def make_text_prompt(args, prompt):
messages = [
get_system_prompt(),
{
"role": "user",
"content": [
{
"type": "text",
"text": prompt
},
]
},
]
prompt = make_inputs_qwen2_omni(args, messages, tokenize=args.tokenize)
return prompt
def make_audio_in_video_v2_prompt(args):
messages = [
{
'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.'
}]
},
{
"role":
"user",
"content": [
{
"type":
"video_url",
"video_url":
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"
},
]
},
]
prompt = make_inputs_qwen2_omni(
args,
messages,
use_audio_in_video=True,
tokenize=args.tokenize,
)
return prompt
def make_omni_prompt(args, prompt = None) -> Union[OmniTokensPrompt, List[OmniTokensPrompt]]:
if args.prompt_type == 'text':
prompt = make_text_prompt(args, prompt)
elif args.prompt_type == 'audio-in-video-v2':
prompt = make_audio_in_video_v2_prompt(args)
else:
raise ValueError(f'Unsupported prompt type: {args.prompt_type}')
return prompt
vllm_omni/config/__init__.py
+29
View File
@@ -1,6 +1,14 @@
"""
Configuration module for vLLM-omni.
"""
from typing import Optional
from pydantic.dataclasses import dataclass
from pydantic import ConfigDict
from vllm.config import ModelConfig
from vllm.config import config
import vllm_omni.model_executor.models as me_models
from .stage_config import (
OmniStageConfig,
@@ -11,7 +19,28 @@ from .stage_config import (
create_dit_stage_config,
)
@config
@dataclass(config=ConfigDict(arbitrary_types_allowed=True))
class OmniModelConfig(ModelConfig):
"""Configuration for Omni models, extending the base ModelConfig."""
stage_id: int = 0
model_stage: str = "thinker"
model_arch: str = "Qwen2_5OmniForConditionalGeneration"
engine_output_type: Optional[str] = None
@property
def registry(self):
return me_models.OmniModelRegistry
@property
def architectures(self) -> list[str]:
return [self.model_arch]
__all__ = [
"OmniModelConfig",
"OmniStageConfig",
"DiTConfig",
"DiTCacheConfig",
vllm_omni/engine/output_processor.py
+397 -201
View File
@@ -1,215 +1,411 @@
"""
Output processing for multimodal outputs in vLLM-omni.
"""
from typing import Dict, Callable, Optional, Any, Union
from typing import List, Dict, Any, Optional, Callable, Union
from vllm.outputs import RequestOutput, CompletionOutput
from vllm.v1.outputs import ModelRunnerOutput as EngineCoreOutput
import torch
from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestState, RequestOutputCollector
from vllm.transformers_utils.tokenizer_group import TokenizerGroup
from vllm.v1.engine import FinishReason
from vllm.v1.metrics.stats import IterationStats
from vllm.sampling_params import RequestOutputKind
from vllm.transformers_utils.tokenizer import AnyTokenizer
from vllm.v1.engine import EngineCoreRequest
from vllm.v1.engine.detokenizer import IncrementalDetokenizer
from vllm.v1.engine.logprobs import LogprobsProcessor
from vllm.v1.engine.parallel_sampling import ParentRequest
from vllm.logger import init_logger
from vllm.v1.engine import EngineCoreOutput
class MultimodalOutputProcessor:
"""Handles multimodal output processing for vLLM-omni."""
def __init__(self):
self.output_handlers: Dict[str, Callable] = {
"image": self._process_image_output,
"text+image": self._process_text_image_output,
"latents": self._process_latents_output,
"text": self._process_text_output,
"pooling": self._process_pooling_output,
}
def process_output(self, engine_core_output: Any) -> List[RequestOutput]:
"""Process engine core output and return formatted RequestOutput."""
if engine_core_output is None:
return []
# If it's already a RequestOutput, return as is
if isinstance(engine_core_output, RequestOutput):
return [engine_core_output]
# If it's a list of RequestOutputs, return as is
if isinstance(engine_core_output, list):
return engine_core_output
# Otherwise, process based on output type
output_type = self._detect_output_type(engine_core_output)
handler = self.output_handlers.get(output_type, self._process_pooling_output)
return handler(engine_core_output)
logger = init_logger(__name__)
def _build_request_output(
class OmniRequestState(RequestState):
def __init__(
self,
source: Any,
completion_outputs: List[CompletionOutput],
) -> RequestOutput:
"""Helper to construct RequestOutput with safe defaults."""
*args,
**kwargs,
):
super().__init__(*args, **kwargs)
self.mm_type: Optional[str] = None
self.mm_accumulated: Optional[torch.Tensor] = None
return RequestOutput(
request_id=getattr(source, "request_id", "unknown"),
prompt=getattr(source, "prompt", ""),
prompt_token_ids=getattr(source, "prompt_token_ids", []),
prompt_logprobs=getattr(source, "prompt_logprobs", None),
outputs=completion_outputs,
finished=getattr(source, "finished", True),
@classmethod
def from_new_request(
cls,
tokenizer: AnyTokenizer,
request: EngineCoreRequest,
prompt: Optional[str],
parent_req: Optional[ParentRequest],
request_index: int,
queue: Optional[Any],
log_stats: bool,
) -> "OmniRequestState":
if sampling_params := request.sampling_params:
if not sampling_params.detokenize:
tokenizer = None
output_kind = sampling_params.output_kind
logprobs_processor = LogprobsProcessor.from_new_request(
tokenizer=tokenizer,
request=request,
)
detokenizer = IncrementalDetokenizer.from_new_request(
tokenizer=tokenizer,
request=request,
)
max_tokens_param = sampling_params.max_tokens
else:
logprobs_processor = None
detokenizer = None
max_tokens_param = None
assert request.pooling_params is not None
output_kind = request.pooling_params.output_kind
return cls(
request_id=request.request_id,
parent_req=parent_req,
request_index=request_index,
lora_name=(request.lora_request.name
if request.lora_request is not None else None),
output_kind=output_kind,
prompt=prompt,
prompt_token_ids=request.prompt_token_ids,
logprobs_processor=logprobs_processor,
detokenizer=detokenizer,
max_tokens_param=max_tokens_param,
arrival_time=request.arrival_time,
queue=queue,
log_stats=log_stats,
)
def _detect_output_type(self, output: Any) -> str:
"""Detect the type of output based on its content."""
if hasattr(output, 'output_type'):
return output.output_type
# Check for image-related attributes
if hasattr(output, 'image') or hasattr(output, 'images'):
if hasattr(output, 'text') or hasattr(output, 'texts'):
return "text+image"
def add_multimodal_tensor(self, tensor: Optional[torch.Tensor],
mm_type: Optional[str]) -> None:
if tensor 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
else:
return "image"
# Check for latent-related attributes
if hasattr(output, 'latents') or hasattr(output, 'latent_representation'):
return "latents"
# Check for pooling output
if hasattr(output, 'pooler_output') and output.pooler_output is not None:
return "pooling"
# Default to text
return "text"
def _process_text_output(self, output: Any) -> List[RequestOutput]:
"""Process text output."""
if isinstance(output, RequestOutput):
return [output]
# Create a mock RequestOutput for text
completion_output = CompletionOutput(
index=0,
text=getattr(output, 'text', ''),
token_ids=getattr(output, 'token_ids', []),
cumulative_logprob=getattr(output, 'cumulative_logprob', 0.0),
logprobs=getattr(output, 'logprobs', None),
finish_reason=getattr(output, 'finish_reason', 'length')
)
return [self._build_request_output(output, [completion_output])]
self.mm_accumulated = torch.cat([self.mm_accumulated, t], dim=0)
except Exception:
# Log and continue without crashing the output pipeline
logger.exception("Error accumulating multimodal tensor")
def _process_image_output(self, output: Any) -> List[RequestOutput]:
"""Process image output."""
# For image outputs, we need to create a special RequestOutput
# that can handle image data
# Override: do not route to pooling-only path; always create completion
# outputs, and attach pooling_result into the CompletionOutput.
def make_request_output(
self,
new_token_ids: list[int],
pooling_output: Optional[torch.Tensor],
finish_reason: Optional[FinishReason],
stop_reason: Optional[Union[int, str]],
kv_transfer_params: Optional[dict[str, Any]] = None,
num_cached_tokens: Optional[int] = None,
) -> Optional[Any]:
finished = finish_reason is not None
final_only = self.output_kind == RequestOutputKind.FINAL_ONLY
if not finished and final_only:
return None
if num_cached_tokens is not None:
# Keep num_cached_tokens in RequestOutput for compatibility
try:
self.num_cached_tokens = num_cached_tokens # type: ignore[attr-defined]
except Exception:
pass
request_id = self.request_id
output = self._new_completion_output(new_token_ids, finish_reason,
stop_reason)
if self.parent_req is None:
outputs = [output]
else:
request_id, outputs, finished = self.parent_req.get_outputs(
request_id, output)
if not outputs:
return None
return self._new_request_output(request_id, outputs, finished,
kv_transfer_params)
def _new_completion_output(
self,
token_ids: list[int],
finish_reason: Optional[FinishReason],
stop_reason: Optional[Union[int, str]]
) -> Any:
# Reuse base text/logprobs logic, then annotate with pooling_result.
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.
if not hasattr(base_output, "multimodal_output"):
setattr(base_output, "multimodal_output", {})
setattr(base_output, "multimodal_output", {self.mm_type: tensor})
except Exception:
logger.exception("Error in _new_completion_output")
return base_output
class MultimodalOutputProcessor(VLLMOutputProcessor):
"""Handles multimodal output processing by normalizing EngineCoreOutput
before delegating to the base vLLM OutputProcessor.
Strategy:
- Route by EngineCoreOutput.output_type when present
("image", "text+image", "latents", "text").
- Fallback to pooling/text heuristics when output_type is absent.
- Mutate EngineCoreOutput in-place to ensure vLLM's base processor can
produce the correct RequestOutput/PoolingRequestOutput.
- Allow custom per-modality handlers via register_handler().
"""
def __init__(self, tokenizer: TokenizerGroup, log_stats: bool, engine_core_output_type: Optional[str] = None):
super().__init__(tokenizer=tokenizer, log_stats=log_stats)
self.output_handlers: Dict[str, Callable[[EngineCoreOutput], None]] = {}
self._reqid_to_mm_type: Dict[str, str] = {}
self.request_states: dict[str, OmniRequestState] = {}
self.engine_core_output_type = engine_core_output_type
# Extract image data
image_data = getattr(output, 'image', None)
if image_data is None:
image_data = getattr(output, 'images', [None])[0]
# Create a completion output with image data
completion_output = CompletionOutput(
index=0,
text="", # No text for pure image output
token_ids=[],
cumulative_logprob=0.0,
logprobs=None,
finish_reason="stop"
def register_handler(self, modality: str,
handler: Callable[[EngineCoreOutput], None]) -> None:
self.output_handlers[modality.lower()] = handler
def add_request(
self,
request: EngineCoreRequest,
prompt: Optional[str],
parent_req: Optional[ParentRequest] = None,
request_index: int = 0,
queue: Optional[RequestOutputCollector] = None,
) -> None:
request_id = request.request_id
if request_id in self.request_states:
raise ValueError(f"Request id {request_id} already running.")
tokenizer = None if not self.tokenizer else \
self.tokenizer.get_lora_tokenizer(request.lora_request)
req_state = OmniRequestState.from_new_request(tokenizer=tokenizer,
request=request,
prompt=prompt,
parent_req=parent_req,
request_index=request_index,
queue=queue,
log_stats=self.log_stats)
self.request_states[request_id] = req_state
self.lora_states.add_request(req_state)
if parent_req:
self.parent_requests[parent_req.request_id] = parent_req
def process_outputs(
self,
engine_core_outputs: list[EngineCoreOutput],
engine_core_timestamp: Optional[float] = None,
iteration_stats: Optional[IterationStats] = None,
) -> OutputProcessorOutput:
self._reqid_to_mm_type.clear()
for eco in engine_core_outputs:
mm_type = (getattr(eco, "output_type", None) or "").lower()
if mm_type:
self._reqid_to_mm_type[eco.request_id] = mm_type
self._route_and_normalize(eco)
# Build RequestOutputs without delegating to base, so we can keep ids
request_outputs: list[Any] = []
reqs_to_abort: list[str] = []
for eco in engine_core_outputs:
req_id = eco.request_id
req_state = self.request_states.get(req_id)
if req_state is None:
continue
# 1) Stats
self._update_stats_from_output(req_state, eco,
engine_core_timestamp,
iteration_stats)
new_token_ids = eco.new_token_ids
pooling_output = eco.pooling_output
finish_reason = eco.finish_reason
stop_reason = eco.stop_reason
kv_transfer_params = eco.kv_transfer_params
num_cached_tokens = eco.num_cached_tokens
req_state.is_prefilling = False
# 2) Detokenize and logprobs when text path
assert req_state.detokenizer is not None
assert req_state.logprobs_processor is not None
stop_string = req_state.detokenizer.update(
new_token_ids, finish_reason == FinishReason.STOP)
if stop_string:
finish_reason = FinishReason.STOP
stop_reason = stop_string
req_state.logprobs_processor.update_from_output(eco)
# 2.5) Accumulate multimodal tensors in RequestState
try:
mm_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
if pooling_output is not None and isinstance(req_state, OmniRequestState):
req_state.add_multimodal_tensor(pooling_output, mm_type)
except Exception:
logger.debug("Failed to accumulate multimodal tensor for request %s", req_id, exc_info=True)
# 3) Create RequestOutput objects, forcing combined mode to keep ids
pooling_for_make = pooling_output
if pooling_output is not None and new_token_ids:
# Do not consume pooling path now; keep ids and attach mm later
pooling_for_make = None
ro = req_state.make_request_output(new_token_ids, pooling_for_make,
finish_reason, stop_reason,
kv_transfer_params,
num_cached_tokens)
if ro:
# 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
except Exception:
logger.exception("Error attaching multimodal payload in process_outputs")
if req_state.queue is not None:
req_state.queue.put(ro)
else:
request_outputs.append(ro)
# 4) Free completed
if finish_reason is not None:
self.request_states.pop(req_id)
parent_req = req_state.parent_req
if parent_req and not parent_req.child_requests:
self.parent_requests.pop(parent_req.request_id, None)
if not eco.finished:
reqs_to_abort.append(req_id)
self._update_stats_from_finished(req_state, finish_reason,
iteration_stats)
# Cleanup per-request mm state
if isinstance(req_state, OmniRequestState):
req_state.mm_accumulated = None
req_state.mm_type = None
return OutputProcessorOutput(
request_outputs=request_outputs,
reqs_to_abort=reqs_to_abort,
)
# Add image data to the completion output
completion_output.image = image_data
return [
self._build_request_output(output, [completion_output])
]
def _process_text_image_output(self, output: Any) -> List[RequestOutput]:
"""Process combined text and image output."""
# Extract text and image data
text_data = getattr(output, 'text', '')
image_data = getattr(output, 'image', None)
if image_data is None:
image_data = getattr(output, 'images', [None])[0]
# Create a completion output with both text and image
completion_output = CompletionOutput(
index=0,
text=text_data,
token_ids=getattr(output, 'token_ids', []),
cumulative_logprob=getattr(output, 'cumulative_logprob', 0.0),
logprobs=getattr(output, 'logprobs', None),
finish_reason="stop"
)
# Add image data to the completion output
completion_output.image = image_data
return [
self._build_request_output(output, [completion_output])
]
def _process_latents_output(self, output: Any) -> List[RequestOutput]:
"""Process latent representation output."""
# Extract latent data
latent_data = getattr(output, 'latents', None)
if latent_data is None:
latent_data = getattr(output, 'latent_representation', None)
# Create a completion output with latent data
completion_output = CompletionOutput(
index=0,
text="", # No text for latent output
token_ids=[],
cumulative_logprob=0.0,
logprobs=None,
finish_reason="stop"
)
# Add latent data to the completion output
completion_output.latents = latent_data
return [
self._build_request_output(output, [completion_output])
]
def _process_pooling_output(self, output: Any) -> List[RequestOutput]:
"""Process pooling output (hidden states, embeddings, etc.)."""
# Extract pooling data
pooling_data = getattr(output, 'pooler_output', None)
if pooling_data is None:
pooling_data = getattr(output, 'hidden_states', None)
# Create a completion output with pooling data
completion_output = CompletionOutput(
index=0,
text="", # No text for pooling output
token_ids=[],
cumulative_logprob=0.0,
logprobs=None,
finish_reason="stop"
)
# Add pooling data to the completion output
completion_output.pooler_output = pooling_data
return [
self._build_request_output(output, [completion_output])
]
def process_outputs(self, engine_core_outputs: List[EngineCoreOutput], **kwargs) -> List[RequestOutput]:
"""Process multiple engine core outputs."""
all_outputs = []
for engine_core_output in engine_core_outputs:
outputs = self.process_output(engine_core_output)
all_outputs.extend(outputs)
return all_outputs
def add_output_handler(self, output_type: str, handler: Callable) -> None:
"""Add a custom output handler for a specific output type."""
self.output_handlers[output_type] = handler
def remove_output_handler(self, output_type: str) -> None:
"""Remove an output handler for a specific output type."""
# ---- routing helpers ----
def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
output_type = (getattr(eco, "output_type", self.engine_core_output_type) or "").lower()
# Custom handler first (if registered)
if output_type in self.output_handlers:
del self.output_handlers[output_type]
try:
self.output_handlers[output_type](eco)
# Fall through to default fixups in case the handler left gaps
except Exception:
logger.exception("Error in custom output handler for %s", output_type)
if output_type == "image":
self._process_image_output(eco)
elif output_type in ("text+image", "text,image", "image+text"):
self._process_text_image_output(eco)
elif output_type in ("latents", "latent"):
self._process_latents_output(eco)
elif output_type in ("audio", "speech"):
self._process_audio_output(eco)
elif output_type == "text":
self._process_text_output(eco)
else:
# Fallback heuristic
if eco.pooling_output is not None:
self._process_pooling_output(eco)
else:
self._process_text_output(eco)
# ---- modality processors ----
def _process_image_output(self, eco: EngineCoreOutput) -> None:
"""Ensure image tensors are surfaced via pooling_output for vLLM."""
if eco.pooling_output is None:
tensor = self._extract_from_multimodal_outputs(
eco, keys=("image", "images", "pixel_values", "pixels"))
if tensor is not None:
eco.pooling_output = tensor
def _process_text_image_output(self, eco: EngineCoreOutput) -> None:
"""Allow text+image outputs. Text path stays as new_token_ids;
image/latents route via pooling_output."""
# Preserve text tokens as-is; ensure pooling_output carries image/latents
if eco.pooling_output is None:
tensor = self._extract_from_multimodal_outputs(
eco, keys=("image", "images", "pixel_values", "pixels",
"latent", "latents", "z"))
if tensor is not None:
eco.pooling_output = tensor
def _process_latents_output(self, eco: EngineCoreOutput) -> None:
"""Ensure latent tensors are surfaced via pooling_output."""
if eco.pooling_output is None:
tensor = self._extract_from_multimodal_outputs(
eco, keys=("latent", "latents", "z", "posterior"))
if tensor is not None:
eco.pooling_output = tensor
def _process_audio_output(self, eco: EngineCoreOutput) -> None:
"""Ensure audio tensors are surfaced via pooling_output."""
if eco.pooling_output is None:
tensor = self._extract_from_multimodal_outputs(
eco, keys=("audio", "audios", "wav", "waveform",
"audio_pcm", "pcm"))
if tensor is not None:
eco.pooling_output = tensor
def _process_text_output(self, eco: EngineCoreOutput) -> None:
"""No-op; base processor will detokenize new_token_ids → text."""
return
def _process_pooling_output(self, eco: EngineCoreOutput) -> None:
"""Optional sanity checks for pooling tensor."""
if eco.pooling_output is None:
return
if not isinstance(eco.pooling_output, torch.Tensor):
# Best-effort: convert to tensor if it's a list/ndarray-like
try:
eco.pooling_output = torch.as_tensor(eco.pooling_output)
except Exception:
pass
def _extract_from_multimodal_outputs(
self, eco: EngineCoreOutput, keys: tuple[str, ...]
) -> Optional[torch.Tensor]:
mm = getattr(eco, "multimodal_outputs", None)
if not isinstance(mm, dict):
return None
for k in keys:
v = mm.get(k)
if isinstance(v, torch.Tensor):
return v
# Try the first tensor in the dict as a fallback
for v in mm.values():
if isinstance(v, torch.Tensor):
return v
return None
vllm_omni/entrypoints/omni_llm.py
+134 -553
View File
@@ -1,576 +1,157 @@
"""
Core OmniLLM and AsyncOmniLLM classes for multi-stage processing.
"""
from typing import Union, Sequence, Optional, Any
import cloudpickle
from pydantic import ValidationError
from typing import List, Dict, Any, Optional, Union, Callable
from vllm.inputs import PromptType
from vllm.sampling_params import SamplingParams
from vllm.entrypoints.llm import LLM
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.engine.llm_engine import LLMEngine
from vllm.outputs import RequestOutput, LoRARequest
from vllm.engine.arg_utils import HfOverrides
from vllm.usage.usage_lib import UsageContext
from vllm.config import CompilationConfig, is_init_field
from vllm.utils import Counter
from vllm.logger import init_logger
import vllm.envs as envs
from ..config import OmniStageConfig
from .stage_manager import StageManager
from ..engine.output_processor import MultimodalOutputProcessor
from ..engine.diffusion_engine import DiffusersPipelineEngine
from vllm_omni.entrypoints.utils import load_stage_configs_from_model
from vllm_omni.entrypoints.omni_stage import OmniStage
from vllm_omni.engine.arg_utils import OmniEngineArgs
from vllm_omni.engine.output_processor import MultimodalOutputProcessor
from vllm_omni.engine.processor import OmniProcessor
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
class OmniLLM(LLM):
"""Extended LLM supporting multiple engines and stage-based processing."""
def __init__(
self,
stage_configs: List[OmniStageConfig],
log_stats: bool = False,
**kwargs
):
# Use the first stage's model as the default model for LLM
default_model = stage_configs[0].model_path if stage_configs else "test-model"
# Track whether we should launch engines in multiprocess mode.
self.multiprocess_mode = kwargs.pop("multiprocess_mode", False)
# Fix configuration validation issues
# Ensure max_num_batched_tokens is at least as large as max_model_len
if 'max_model_len' in kwargs and 'max_num_batched_tokens' in kwargs:
if kwargs['max_num_batched_tokens'] < kwargs['max_model_len']:
kwargs['max_num_batched_tokens'] = kwargs['max_model_len']
elif 'max_model_len' in kwargs:
# If max_model_len is set but max_num_batched_tokens is not, set it to max_model_len
kwargs['max_num_batched_tokens'] = kwargs['max_model_len']
class OmniLLM:
def __init__(self, model: str, stage_configs = None, log_stats: bool = False, **kwargs):
if stage_configs is None:
self.initialize_stage_configs(model)
else:
# Set reasonable defaults to avoid validation errors
kwargs['max_model_len'] = 2048
kwargs['max_num_batched_tokens'] = 2048
self.stage_configs = stage_configs
super().__init__(model=default_model, **kwargs)
self.stage_configs = stage_configs
self.log_stats = log_stats
self.stage_manager = StageManager(stage_configs, log_stats)
self.output_processor = MultimodalOutputProcessor()
self._initialize_stage_engines()
self.stage_list = []
self.initialize_stages(model)
def initialize_stage_configs(self, model: str):
self.stage_configs = load_stage_configs_from_model(model)
def _initialize_stage_engines(self) -> None:
"""Initialize LLMEngine instances for each stage."""
self.stage_manager.initialize_engines()
def initialize_stages(self, model: str):
for stage_config in self.stage_configs:
stage = OmniStage(stage_config)
stage_llm = OmniStageLLM(model=model, **stage_config.engine_args)
stage.set_engine(stage_llm)
self.stage_list.append(stage)
def generate(
self,
stage_args_list: Optional[List[Dict[str, Any]]] = None,
use_tqdm: Union[bool, Callable[..., Any]] = True,
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
priority: Optional[List[int]] = None,
*,
prompt: Optional[str] = None,
stage_overrides: Optional[Dict[int, Dict[str, Any]]] = None,
) -> List[RequestOutput]:
"""Main generation interface - orchestrates multi-stage processing."""
if stage_args_list is None:
if prompt is None:
prompts: Union[PromptType, Sequence[PromptType]],
sampling_params_list: Optional[Union[SamplingParams,
Sequence[SamplingParams]]] = None,
) -> list[OmniRequestOutput]:
"""Generate text outputs for the given prompts."""
final_outputs: list[OmniRequestOutput] = []
if len(sampling_params_list) != len(self.stage_list):
raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
for stage_id, stage in enumerate(self.stage_list):
if stage_id > 0:
engine_inputs = stage.process_engine_inputs(self.stage_list, prompts)
else:
engine_inputs = prompts
engine_outputs = self._run_generation(stage, sampling_params_list[stage_id], engine_inputs)
stage.set_engine_outputs(engine_outputs)
if hasattr(stage, 'final_output') and stage.final_output:
final_outputs.append(OmniRequestOutput(
stage_id=stage_id,
final_output_type=stage.final_output_type,
request_output=engine_outputs))
return final_outputs
def _run_generation(self, stage: OmniStage, sampling_params: SamplingParams, prompts: Union[PromptType, Sequence[PromptType]]):
engine_outputs = []
for ro in stage.engine.generate(prompts, sampling_params):
engine_outputs.append(ro)
return engine_outputs
class OmniStageLLM(LLM):
def __init__(self,
model: str,
compilation_config: Optional[Union[int, dict[str, Any],
CompilationConfig]] = None,
hf_overrides: Optional[HfOverrides] = None,
**kwargs):
"""LLM constructor."""
if "disable_log_stats" not in kwargs:
kwargs["disable_log_stats"] = True
if "worker_cls" in kwargs:
worker_cls = kwargs["worker_cls"]
# if the worker_cls is not qualified string name,
# we serialize it using cloudpickle to avoid pickling issues
if isinstance(worker_cls, type):
kwargs["worker_cls"] = cloudpickle.dumps(worker_cls)
if "kv_transfer_config" in kwargs and isinstance(
kwargs["kv_transfer_config"], dict):
from vllm.config import KVTransferConfig
raw_config_dict = kwargs["kv_transfer_config"]
try:
kwargs["kv_transfer_config"] = KVTransferConfig(
**raw_config_dict)
except ValidationError as e:
logger.error(
"Failed to convert 'kv_transfer_config' dict to "
"KVTransferConfig object. Dict: %s. Error: %s",
raw_config_dict, e)
# Consider re-raising a more specific vLLM error or ValueError
# to provide better context to the user.
raise ValueError(
"prompt must be provided when stage_args_list is not supplied"
)
stage_args_list = self._build_stage_args_from_config(
prompt, stage_overrides or {}
)
f"Invalid 'kv_transfer_config' provided: {e}") from e
if len(stage_args_list) != len(self.stage_configs):
raise ValueError(
f"Number of stage arguments ({len(stage_args_list)}) must match "
f"number of stage configs ({len(self.stage_configs)})"
)
# Process through each stage sequentially
current_output = None
for i, (stage_config, stage_args) in enumerate(zip(self.stage_configs, stage_args_list)):
stage_engine = self.stage_manager.get_engine(i)
# Prepare input for this stage
processed_input = self._process_stage_inputs(
stage_config, stage_args or {}, current_output
)
# Execute stage
stage_output = self._execute_stage(
stage_engine, processed_input, lora_request, priority, stage_config
)
# Update for next stage
current_output = stage_output
stage_config.stage_output = stage_output
# Process final output
final_output = self.output_processor.process_output(current_output)
return final_output
def _build_stage_args_from_config(
self,
prompt: str,
stage_overrides: Dict[int, Dict[str, Any]],
) -> List[Dict[str, Any]]:
"""Derive per-stage argument dictionaries from configuration defaults."""
stage_args: List[Dict[str, Any]] = []
for stage_config in self.stage_configs:
combined: Dict[str, Any] = dict(stage_config.default_stage_args or {})
override = stage_overrides.get(stage_config.stage_id)
if override:
combined.update(override)
if stage_config.engine_type == "AR":
combined["prompt"] = prompt
stage_args.append(combined)
return stage_args
if hf_overrides is None:
hf_overrides = {}
def _process_stage_inputs(
self,
stage_config: OmniStageConfig,
stage_args: Dict[str, Any],
previous_output: Optional[Any]
) -> Dict[str, Any]:
"""Prepare input for specific stage."""
if stage_config.engine_type == "AR":
return self._process_ar_inputs(stage_config, stage_args, previous_output)
elif stage_config.engine_type == "DiT":
return self._process_dit_inputs(stage_config, stage_args, previous_output)
if compilation_config is not None:
if isinstance(compilation_config, int):
compilation_config_instance = CompilationConfig(
level=compilation_config)
elif isinstance(compilation_config, dict):
predicate = lambda x: is_init_field(CompilationConfig, x[0])
compilation_config_instance = CompilationConfig(
**dict(filter(predicate, compilation_config.items())))
else:
compilation_config_instance = compilation_config
else:
raise NotImplementedError(f"Unknown engine type: {stage_config.engine_type}")
compilation_config_instance = CompilationConfig()
def _process_ar_inputs(
self,
stage_config: OmniStageConfig,
stage_args: Dict[str, Any],
previous_output: Optional[Any]
) -> Dict[str, Any]:
"""Process inputs for AR stage."""
combined = dict(stage_config.default_stage_args or {})
combined.update(stage_args)
combined.setdefault("prompt", "")
combined.setdefault("max_tokens", 100)
combined.setdefault("temperature", 0.7)
# If we have previous output (e.g., from a previous AR stage),
# we might want to use it as context
if previous_output is not None:
# Extract text from previous output if available
if hasattr(previous_output, 'outputs') and previous_output.outputs:
last_output = previous_output.outputs[-1]
if hasattr(last_output, 'text'):
combined["prompt"] = last_output.text + " " + combined["prompt"]
return combined
def _process_dit_inputs(
self,
stage_config: OmniStageConfig,
stage_args: Dict[str, Any],
previous_output: Optional[Any]
) -> Dict[str, Any]:
"""Process inputs for DiT stage."""
combined = dict(stage_config.default_stage_args or {})
combined.update(stage_args)
dit = stage_config.dit_config
if dit is not None:
combined.setdefault("height", getattr(dit, "height", 512))
combined.setdefault("width", getattr(dit, "width", 512))
combined.setdefault(
"num_inference_steps", getattr(dit, "num_inference_steps", 50)
)
combined.setdefault(
"guidance_scale", getattr(dit, "guidance_scale", 7.5)
)
else:
combined.setdefault("height", 512)
combined.setdefault("width", 512)
combined.setdefault("num_inference_steps", 50)
combined.setdefault("guidance_scale", 7.5)
# Handle image inputs if present
if "image" in stage_args:
# For now, we'll pass the image path directly
# In a full implementation, this would involve VAE encoding
combined["image"] = stage_args["image"]
# If we have previous output from an AR stage, we might want to use it
if previous_output is not None:
# Extract text from previous AR output
if hasattr(previous_output, 'outputs') and previous_output.outputs:
last_output = previous_output.outputs[-1]
if hasattr(last_output, 'text'):
combined["prompt"] = last_output.text
combined.setdefault("prompt", stage_args.get("prompt", ""))
return combined
def _execute_stage(
self,
stage_engine: Optional[Union[LLMEngine, DiffusersPipelineEngine]],
processed_input: Dict[str, Any],
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
priority: Optional[List[int]] = None,
stage_config: Optional[OmniStageConfig] = None,
) -> Any:
"""Execute a single stage."""
# DiT via diffusers backend
if stage_config and stage_config.engine_type == "DiT":
dit = stage_config.dit_config
if dit and getattr(dit, "use_diffusers", False):
# Lazy-init executor per stage
if not hasattr(self, "_dit_engines"):
self._dit_engines = {}
exec_inst = self._dit_engines.get(stage_config.stage_id)
if exec_inst is None:
exec_inst = DiffusersPipelineEngine(
dit_config=dit,
model_path=stage_config.model_path,
log_stats=self.log_stats,
multiprocess_mode=self.multiprocess_mode,
)
self._dit_engines[stage_config.stage_id] = exec_inst
return exec_inst.generate(
prompt=processed_input.get("prompt", ""),
height=processed_input.get("height", getattr(dit, "height", 512)),
width=processed_input.get("width", getattr(dit, "width", 512)),
num_inference_steps=processed_input.get(
"num_inference_steps", getattr(dit, "num_inference_steps", 30)
),
guidance_scale=processed_input.get(
"guidance_scale", getattr(dit, "guidance_scale", 5.0)
),
negative_prompt=processed_input.get("negative_prompt"),
seed=processed_input.get("seed"),
image=processed_input.get("image"),
)
# Use the parent LLM's generate method for AR text generation
prompt = processed_input.get("prompt", "")
max_tokens = processed_input.get("max_tokens", 100)
temperature = processed_input.get("temperature", 0.7)
# Generate using the base LLM class
from vllm.sampling_params import SamplingParams
sampling_params = SamplingParams(
max_tokens=max_tokens,
temperature=temperature,
top_p=processed_input.get("top_p", 1.0),
frequency_penalty=processed_input.get("frequency_penalty", 0.0),
presence_penalty=processed_input.get("presence_penalty", 0.0),
stop=processed_input.get("stop", None)
engine_args = OmniEngineArgs(
model=model,
hf_overrides=hf_overrides,
compilation_config=compilation_config_instance,
**kwargs,
)
# Use the parent class's generate method
outputs = super().generate([prompt], sampling_params)
# Return the first output (we're processing one prompt at a time)
if outputs:
return outputs[0]
# Create the Engine (autoselects V0 vs V1)
self.llm_engine = LLMEngine.from_engine_args(
engine_args=engine_args, usage_context=UsageContext.LLM_CLASS)
self.llm_engine.output_processor = MultimodalOutputProcessor(tokenizer=self.llm_engine.tokenizer,
log_stats=self.llm_engine.log_stats)
self.llm_engine.processor = OmniProcessor(vllm_config=self.llm_engine.vllm_config,
tokenizer=self.llm_engine.tokenizer)
self.engine_class = type(self.llm_engine)
self.request_counter = Counter()
self.default_sampling_params: Union[dict[str, Any], None] = None
if envs.VLLM_USE_V1:
supported_tasks = self.llm_engine \
.get_supported_tasks() # type: ignore
else:
# Fallback to mock output if generation fails
from vllm.outputs import RequestOutput, CompletionOutput
mock_output = CompletionOutput(
index=0,
text="Generation failed",
token_ids=[],
cumulative_logprob=0.0,
logprobs=None,
finish_reason="error"
)
return RequestOutput(
request_id="fallback_request",
prompt=prompt,
prompt_token_ids=[],
prompt_logprobs=None,
outputs=[mock_output],
finished=True
)
supported_tasks = self.llm_engine.model_config.supported_tasks
logger.info("Supported_tasks: %s", supported_tasks)
class AsyncOmniLLM(LLM):
"""Extended LLM class supporting multiple engines and stage-based processing."""
def __init__(
self,
stage_configs: List[OmniStageConfig],
log_stats: bool = False,
**kwargs
):
# Use the first stage's model for the base LLM
if stage_configs and stage_configs[0].model_path:
model = stage_configs[0].model_path
else:
model = "Qwen/Qwen3-0.6B"
# Fix configuration validation issues
# Ensure max_num_batched_tokens is at least as large as max_model_len
if 'max_model_len' in kwargs and 'max_num_batched_tokens' in kwargs:
if kwargs['max_num_batched_tokens'] < kwargs['max_model_len']:
kwargs['max_num_batched_tokens'] = kwargs['max_model_len']
elif 'max_model_len' in kwargs:
# If max_model_len is set but max_num_batched_tokens is not, set it to max_model_len
kwargs['max_num_batched_tokens'] = kwargs['max_model_len']
else:
# Set reasonable defaults to avoid validation errors
kwargs['max_model_len'] = 2048
kwargs['max_num_batched_tokens'] = 2048
super().__init__(model=model, **kwargs)
self.stage_configs = stage_configs
self.log_stats = log_stats
self.stage_manager = StageManager(stage_configs, log_stats)
self.output_processor = MultimodalOutputProcessor()
def _initialize_async_stage_engines(self) -> None:
"""Initialize AsyncLLM instances for each stage."""
self.stage_manager.initialize_async_engines()
async def generate_async(
self,
stage_args_list: Optional[List[Dict[str, Any]]] = None,
use_tqdm: Union[bool, Callable[..., Any]] = True,
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
priority: Optional[List[int]] = None,
*,
prompt: Optional[str] = None,
stage_overrides: Optional[Dict[int, Dict[str, Any]]] = None,
) -> List[RequestOutput]:
"""Async generation interface - orchestrates multi-stage processing."""
if stage_args_list is None:
if prompt is None:
raise ValueError(
"prompt must be provided when stage_args_list is not supplied"
)
stage_args_list = self._build_stage_args_from_config(
prompt, stage_overrides or {}
)
if len(stage_args_list) != len(self.stage_configs):
raise ValueError(
f"Number of stage arguments ({len(stage_args_list)}) must match "
f"number of stage configs ({len(self.stage_configs)})"
)
# Process through each stage sequentially
current_output = None
for i, (stage_config, stage_args) in enumerate(zip(self.stage_configs, stage_args_list)):
stage_engine = self.stage_manager.get_async_engine(i)
# Prepare input for this stage
processed_input = self._process_stage_inputs(
stage_config, stage_args or {}, current_output
)
# Execute stage asynchronously
stage_output = await self._execute_stage_async(
stage_engine, processed_input, lora_request, priority, stage_config
)
# Update for next stage
current_output = stage_output
stage_config.stage_output = stage_output
# Process final output
final_output = self.output_processor.process_output(current_output)
return final_output
def _build_stage_args_from_config(
self,
prompt: str,
stage_overrides: Dict[int, Dict[str, Any]],
) -> List[Dict[str, Any]]:
stage_args: List[Dict[str, Any]] = []
for stage_config in self.stage_configs:
combined: Dict[str, Any] = dict(stage_config.default_stage_args or {})
override = stage_overrides.get(stage_config.stage_id)
if override:
combined.update(override)
if stage_config.engine_type == "AR":
combined["prompt"] = prompt
stage_args.append(combined)
return stage_args
def _process_stage_inputs(
self,
stage_config: OmniStageConfig,
stage_args: Dict[str, Any],
previous_output: Optional[Any]
) -> Dict[str, Any]:
"""Prepare input for specific stage (same as OmniLLM)."""
if stage_config.engine_type == "AR":
return self._process_ar_inputs(stage_config, stage_args, previous_output)
elif stage_config.engine_type == "DiT":
return self._process_dit_inputs(stage_config, stage_args, previous_output)
else:
raise NotImplementedError(f"Unknown engine type: {stage_config.engine_type}")
def _process_ar_inputs(
self,
stage_config: OmniStageConfig,
stage_args: Dict[str, Any],
previous_output: Optional[Any]
) -> Dict[str, Any]:
"""Process inputs for AR stage (same as OmniLLM)."""
combined = dict(stage_config.default_stage_args or {})
combined.update(stage_args)
combined.setdefault("prompt", "")
combined.setdefault("max_tokens", 100)
combined.setdefault("temperature", 0.7)
if previous_output is not None:
if hasattr(previous_output, 'outputs') and previous_output.outputs:
last_output = previous_output.outputs[-1]
if hasattr(last_output, 'text'):
combined["prompt"] = last_output.text + " " + combined["prompt"]
return combined
def _process_dit_inputs(
self,
stage_config: OmniStageConfig,
stage_args: Dict[str, Any],
previous_output: Optional[Any]
) -> Dict[str, Any]:
"""Process inputs for DiT stage (same as OmniLLM)."""
combined = dict(stage_config.default_stage_args or {})
combined.update(stage_args)
dit = stage_config.dit_config
if dit is not None:
combined.setdefault("height", getattr(dit, "height", 512))
combined.setdefault("width", getattr(dit, "width", 512))
combined.setdefault(
"num_inference_steps", getattr(dit, "num_inference_steps", 50)
)
combined.setdefault(
"guidance_scale", getattr(dit, "guidance_scale", 7.5)
)
else:
combined.setdefault("height", 512)
combined.setdefault("width", 512)
combined.setdefault("num_inference_steps", 50)
combined.setdefault("guidance_scale", 7.5)
if "image" in stage_args:
combined["image"] = stage_args["image"]
if previous_output is not None:
if hasattr(previous_output, 'outputs') and previous_output.outputs:
last_output = previous_output.outputs[-1]
if hasattr(last_output, 'text'):
combined["prompt"] = last_output.text
combined.setdefault("prompt", stage_args.get("prompt", ""))
return combined
async def _execute_stage_async(
self,
stage_engine: AsyncLLM,
processed_input: Dict[str, Any],
lora_request: Optional[Union[List[LoRARequest], LoRARequest]] = None,
priority: Optional[List[int]] = None,
stage_config: Optional[OmniStageConfig] = None,
) -> Any:
"""Execute a single stage asynchronously."""
# DiT via diffusers backend (sync call inside async for MVP)
if stage_config and stage_config.engine_type == "DiT":
dit = stage_config.dit_config
if dit and getattr(dit, "use_diffusers", False):
if not hasattr(self, "_dit_engines"):
self._dit_engines = {}
exec_inst = self._dit_engines.get(stage_config.stage_id)
if exec_inst is None:
from vllm_omni.engine.diffusion_engine import (
DiffusersPipelineEngine,
)
pipeline_name = getattr(dit, "diffusers_pipeline", None)
device_cfg = getattr(dit, "device_config", None)
model_cfg = getattr(dit, "model_config", None)
if isinstance(device_cfg, dict):
device = device_cfg.get("device")
dtype = device_cfg.get("dtype")
else:
device = getattr(device_cfg, "device", None)
dtype = getattr(device_cfg, "dtype", None)
if dtype is None:
if isinstance(model_cfg, dict):
dtype = model_cfg.get("dtype")
else:
dtype = getattr(model_cfg, "dtype", None)
exec_inst = DiffusersPipelineEngine(
model_path=stage_config.model_path,
pipeline_name=pipeline_name,
device=device,
dtype=dtype,
)
self._dit_engines[stage_config.stage_id] = exec_inst
return exec_inst.generate(
prompt=processed_input.get("prompt", ""),
height=processed_input.get("height", getattr(dit, "height", 512)),
width=processed_input.get("width", getattr(dit, "width", 512)),
num_inference_steps=processed_input.get(
"num_inference_steps", getattr(dit, "num_inference_steps", 30)
),
guidance_scale=processed_input.get(
"guidance_scale", getattr(dit, "guidance_scale", 5.0)
),
negative_prompt=processed_input.get("negative_prompt"),
seed=processed_input.get("seed"),
image=processed_input.get("image"),
)
# Use the parent LLM's generate method for AR text generation
prompt = processed_input.get("prompt", "")
max_tokens = processed_input.get("max_tokens", 100)
temperature = processed_input.get("temperature", 0.7)
# Generate using the base LLM class
from vllm.sampling_params import SamplingParams
sampling_params = SamplingParams(
max_tokens=max_tokens,
temperature=temperature,
top_p=processed_input.get("top_p", 1.0),
frequency_penalty=processed_input.get("frequency_penalty", 0.0),
presence_penalty=processed_input.get("presence_penalty", 0.0),
stop=processed_input.get("stop", None)
)
# Use the parent class's generate method
outputs = super().generate([prompt], sampling_params)
# Return the first output (we're processing one prompt at a time)
if outputs:
return outputs[0]
else:
# Fallback to mock output if generation fails
from vllm.outputs import RequestOutput, CompletionOutput
mock_output = CompletionOutput(
index=0,
text="Generation failed",
token_ids=[],
cumulative_logprob=0.0,
logprobs=None,
finish_reason="error"
)
return RequestOutput(
request_id="fallback_request",
prompt=prompt,
prompt_token_ids=[],
prompt_logprobs=None,
outputs=[mock_output],
finished=True
)
self.supported_tasks = supported_tasks
vllm_omni/entrypoints/omni_stage.py
+80
View File
@@ -0,0 +1,80 @@
"""
Stage manager for orchestrating multiple engines in vLLM-omni.
"""
import importlib
from typing import List, Union
from vllm.v1.engine.llm_engine import LLMEngine
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.inputs import TextPrompt
from vllm.v1.engine import EngineCoreOutput
from vllm_omni.inputs.data import OmniTokensPrompt
class OmniStage:
def __init__(self, stage_config):
self.stage_config = stage_config
self.engine = None
self.async_engine = None
self.stage_id = stage_config.stage_id
self.engine_args = stage_config.engine_args
self.model_stage = stage_config.engine_args.model_stage
if hasattr(stage_config, 'engine_input_source'):
self.engine_input_source = stage_config.engine_input_source
else:
self.engine_input_source = []
self.engine_output_type = stage_config.engine_args.engine_output_type
self.engine_outputs = None
if hasattr(stage_config, 'custom_process_input_func'):
# Import the module specified in the config (already a full module path)
module_path, func_name = stage_config.custom_process_input_func.rsplit('.', 1)
module = importlib.import_module(module_path)
self.custom_process_input_func = getattr(module, func_name)
else:
self.custom_process_input_func = None
if hasattr(stage_config, 'final_output'):
self.final_output = stage_config.final_output
else:
self.final_output = False
if hasattr(stage_config, 'final_output_type'):
self.final_output_type = stage_config.final_output_type
else:
self.final_output_type = None
def set_engine(self, engine: LLMEngine) -> None:
"""Initialize the engine for the stage."""
self.engine = engine
def set_async_engine(self, async_engine: AsyncLLM) -> None:
"""Initialize the async engine for the stage."""
self.async_engine = async_engine
def set_engine_outputs(self, engine_outputs: EngineCoreOutput) -> None:
"""Set the engine output for the stage."""
self.engine_outputs = engine_outputs
def process_engine_inputs(self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None) -> List[Union[OmniTokensPrompt, TextPrompt]]:
"""Process the engine input for the stage."""
if self.custom_process_input_func is None:
engine_inputs = []
if len(self.engine_input_source) == 0:
raise ValueError("engine_input_source is empty")
source_stage_id = self.engine_input_source[0]
source_outputs = stage_list[source_stage_id].engine_outputs
multi_modal_data = {source_output.request_id:
prompt.get('multi_modal_data', None) for source_output, prompt in zip(source_outputs, prompt)}
for source_output in source_outputs:
engine_input = OmniTokensPrompt(
prompt_token_ids = source_output.outputs[0].token_ids,
multi_modal_data=multi_modal_data[source_output.request_id] if multi_modal_data else None,
)
engine_inputs.append(engine_input)
return engine_inputs
else:
engine_input_source = self.engine_input_source
return self.custom_process_input_func(stage_list, engine_input_source, prompt)
vllm_omni/entrypoints/stage_manager.py
-87
View File
@@ -1,87 +0,0 @@
"""
Stage manager for orchestrating multiple engines in vLLM-omni.
"""
from typing import List, Optional, Union
from vllm.v1.engine.llm_engine import LLMEngine
from vllm.v1.engine.async_llm import AsyncLLM
from ..engine.diffusion_engine import DiffusionEngine
from ..config import OmniStageConfig
class StageManager:
"""Manages multiple stage engines for multi-stage processing."""
def __init__(self, stage_configs: List[OmniStageConfig], log_stats: bool = False):
self.stage_configs = stage_configs
self.log_stats = log_stats
self.engine_list: List[Union[LLMEngine, DiffusionEngine]] = []
self.async_engine_list: List[AsyncLLM] = []
self._initialized = False
self._async_initialized = False
def initialize_engines(self) -> None:
"""Initialize LLMEngine instances for each stage."""
if self._initialized:
return
# For now, create placeholder engines
# In a full implementation, this would create actual engines
for stage_config in self.stage_configs:
# Placeholder - would create actual engine here
self.engine_list.append(None)
self._initialized = True
def initialize_async_engines(self) -> None:
"""Initialize AsyncLLM instances for each stage."""
if self._async_initialized:
return
# For now, create placeholder engines
# In a full implementation, this would create actual engines
for stage_config in self.stage_configs:
# Placeholder - would create actual engine here
self.async_engine_list.append(None)
self._async_initialized = True
def get_engine(self, stage_id: int) -> LLMEngine:
"""Get the engine for a specific stage."""
if not self._initialized:
self.initialize_engines()
if stage_id >= len(self.engine_list):
raise IndexError(f"Stage {stage_id} not found. Available stages: 0-{len(self.engine_list)-1}")
return self.engine_list[stage_id]
def get_async_engine(self, stage_id: int) -> AsyncLLM:
"""Get the async engine for a specific stage."""
if not self._async_initialized:
self.initialize_async_engines()
if stage_id >= len(self.async_engine_list):
raise IndexError(f"Async stage {stage_id} not found. Available stages: 0-{len(self.async_engine_list)-1}")
return self.async_engine_list[stage_id]
def get_stage_config(self, stage_id: int) -> OmniStageConfig:
"""Get the configuration for a specific stage."""
if stage_id >= len(self.stage_configs):
raise IndexError(f"Stage config {stage_id} not found. Available stages: 0-{len(self.stage_configs)-1}")
return self.stage_configs[stage_id]
def get_num_stages(self) -> int:
"""Get the number of stages."""
return len(self.stage_configs)
def cleanup(self) -> None:
"""Clean up resources."""
# Clean up engines if needed
self.engine_list.clear()
self.async_engine_list.clear()
self._initialized = False
self._async_initialized = False
vllm_omni/entrypoints/utils.py
+25
View File
@@ -0,0 +1,25 @@
import os
from pathlib import Path
from omegaconf import OmegaConf
from vllm.transformers_utils.config import get_config
# Get the project root directory (2 levels up from this file)
PROJECT_ROOT = Path(__file__).parent.parent.parent
def load_stage_configs_from_model(model: str):
"""Load stage configs from model."""
hf_config = get_config(model, trust_remote_code=True)
model_type = hf_config.model_type
stage_config_file = f"vllm_omni/model_executor/stage_configs/{model_type}.yaml"
stage_config_path = PROJECT_ROOT / stage_config_file
if not os.path.exists(stage_config_path):
raise FileNotFoundError(f"Stage config file {stage_config_path} not found")
stage_configs = load_stage_configs_from_yaml(config_path=str(stage_config_path))
return stage_configs
def load_stage_configs_from_yaml(config_path: str):
"""Load stage configs from yaml file."""
config_data = OmegaConf.load(config_path)
return config_data.stage_args
vllm_omni/model_executor/layers/__init__.py
View File
Whitespace-only changes.
vllm_omni/model_executor/layers/mrope.py
+699
View File
@@ -0,0 +1,699 @@
import itertools
from typing import Optional, Union
import numpy as np
import torch
from transformers import PretrainedConfig
from vllm.model_executor.layers.rotary_embedding.base import RotaryEmbedding
from vllm.model_executor.layers.rotary_embedding.common import apply_rotary_emb_dispatch
def _apply_rotary_emb(
x: torch.Tensor,
cos: torch.Tensor,
sin: torch.Tensor,
is_neox_style: bool,
) -> torch.Tensor:
"""
Args:
x: [num_tokens, num_heads, head_size]
cos: [num_tokens, head_size // 2]
sin: [num_tokens, head_size // 2]
is_neox_style: Whether to use the Neox-style or GPT-J-style rotary
positional embeddings.
"""
cos = cos.unsqueeze(-2).to(x.dtype)
sin = sin.unsqueeze(-2).to(x.dtype)
if is_neox_style:
x1, x2 = torch.chunk(x, 2, dim=-1)
else:
x1 = x[..., ::2]
x2 = x[..., 1::2]
o1 = x1 * cos - x2 * sin
o2 = x2 * cos + x1 * sin
if is_neox_style:
return torch.cat((o1, o2), dim=-1)
else:
return torch.stack((o1, o2), dim=-1).flatten(-2)
class MRotaryEmbedding(RotaryEmbedding):
"""Rotary Embedding with Multimodal Sections."""
def __init__(
self,
head_size: int,
rotary_dim: int,
max_position_embeddings: int,
base: float,
is_neox_style: bool,
dtype: torch.dtype,
mrope_section: Optional[list[int]] = None,
) -> None:
# In Qwen2.5-VL, the maximum index value is related to the duration of
# the input video. We enlarge max_position_embeddings to 4 times to get
# a larger the cos and sin cache.
self.cache_max_position_num = max_position_embeddings * 4
super().__init__(head_size, rotary_dim, self.cache_max_position_num,
base, is_neox_style, dtype)
self.mrope_section = mrope_section
if self.mrope_section:
print("Warning: mrope_section check is disabled in Qwen2.5-Omni, this may cause errors, and should be restored in the future.")
# assert sum(self.mrope_section) == rotary_dim // 2
def forward(
self,
positions: torch.Tensor,
query: torch.Tensor,
key: Optional[torch.Tensor] = None,
) -> tuple[torch.Tensor, Optional[torch.Tensor]]:
"""PyTorch-native implementation equivalent to forward().
Args:
positions:
[num_tokens,] (text only) or
[3, num_tokens] (T/H/W positions with multimodal inputs)
query: [num_tokens, num_heads * head_size]
key: [num_tokens, num_kv_heads * head_size]
"""
assert positions.ndim == 1 or positions.ndim == 2
assert key is not None
num_tokens = positions.shape[-1]
cos_sin = self.cos_sin_cache[positions]
cos, sin = cos_sin.chunk(2, dim=-1)
if positions.ndim == 2:
assert self.mrope_section
cos = torch.cat([
m[i]
for i, m in enumerate(cos.split(self.mrope_section, dim=-1))
],
dim=-1)
sin = torch.cat([
m[i]
for i, m in enumerate(sin.split(self.mrope_section, dim=-1))
],
dim=-1)
query_shape = query.shape
query = query.view(num_tokens, -1, self.head_size)
query_rot = query[..., :self.rotary_dim]
query_pass = query[..., self.rotary_dim:]
# query_rot = apply_rotary_emb_dispatch(query_rot, cos, sin,
# self.is_neox_style)
query_rot = _apply_rotary_emb(query_rot, cos, sin,
self.is_neox_style)
query = torch.cat((query_rot, query_pass), dim=-1).reshape(query_shape)
key_shape = key.shape
key = key.view(num_tokens, -1, self.head_size)
key_rot = key[..., :self.rotary_dim]
key_pass = key[..., self.rotary_dim:]
# key_rot = apply_rotary_emb_dispatch(key_rot, cos, sin,
# self.is_neox_style)
key_rot = _apply_rotary_emb(key_rot, cos, sin,
self.is_neox_style)
key = torch.cat((key_rot, key_pass), dim=-1).reshape(key_shape)
return query, key
@classmethod
def get_input_positions(
cls,
input_tokens: list[int],
hf_config: PretrainedConfig,
image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
second_per_grid_ts: Optional[list[float]],
context_len: int = 0,
seq_len: Optional[int] = None,
audio_feature_lengths: Optional[torch.Tensor] = None,
use_audio_in_video: bool = False,
) -> tuple[list[list[int]], int]:
"""Get mrope input positions and delta value."""
image_grid_thw = [] if image_grid_thw is None else image_grid_thw
video_grid_thw = [] if video_grid_thw is None else video_grid_thw
second_per_grid_ts = [] if second_per_grid_ts is None else \
second_per_grid_ts
llm_positions, mrope_position_delta = \
cls.get_input_positions_tensor(
input_tokens=input_tokens,
hf_config=hf_config,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
second_per_grid_ts=second_per_grid_ts,
context_len=context_len,
seq_len=seq_len,
audio_feature_lengths=audio_feature_lengths,
use_audio_in_video=use_audio_in_video,
)
return llm_positions.tolist(), mrope_position_delta
@classmethod
def get_input_positions_tensor(
cls,
input_tokens: list[int],
hf_config: PretrainedConfig,
image_grid_thw: Union[list[list[int]], torch.Tensor],
video_grid_thw: Union[list[list[int]], torch.Tensor],
second_per_grid_ts: list[float],
context_len: int = 0,
seq_len: Optional[int] = None,
audio_feature_lengths: Optional[torch.Tensor] = None,
use_audio_in_video: bool = False,
) -> tuple[torch.Tensor, int]:
from vllm.transformers_utils.config import thinker_uses_mrope
if thinker_uses_mrope(hf_config):
return cls._omni_get_input_positions_tensor(
input_tokens=input_tokens,
hf_config=hf_config,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
second_per_grid_ts=second_per_grid_ts,
context_len=context_len,
seq_len=seq_len,
audio_feature_lengths=audio_feature_lengths,
use_audio_in_video=use_audio_in_video,
)
elif hf_config.model_type in ["glm4v", "glm4v_moe"]:
return cls._glm4v_get_input_positions_tensor(
input_tokens=input_tokens,
hf_config=hf_config,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
context_len=context_len,
seq_len=seq_len,
)
else:
return cls._vl_get_input_positions_tensor(
input_tokens=input_tokens,
hf_config=hf_config,
image_grid_thw=image_grid_thw,
video_grid_thw=video_grid_thw,
second_per_grid_ts=second_per_grid_ts,
context_len=context_len,
seq_len=seq_len,
)
@classmethod
def _glm4v_get_input_positions_tensor(
cls,
input_tokens: list[int],
hf_config: PretrainedConfig,
image_grid_thw: Union[list[list[int]], torch.Tensor],
video_grid_thw: Union[list[list[int]], torch.Tensor],
context_len: int = 0,
seq_len: Optional[int] = None,
) -> tuple[torch.Tensor, int]:
"""Get mrope input positions and delta value for GLM4V."""
image_token_id = hf_config.image_token_id
video_start_token_id = hf_config.video_start_token_id
video_end_token_id = hf_config.video_end_token_id
spatial_merge_size = hf_config.vision_config.spatial_merge_size
llm_pos_ids_list: list = []
if not (image_grid_thw is None and video_grid_thw is None):
if isinstance(image_grid_thw, torch.Tensor):
image_grid_thw = image_grid_thw.tolist()
input_token_type: list[str] = []
video_check_flg = False
for token in input_tokens:
if token == video_start_token_id:
video_check_flg = True
elif token == video_end_token_id:
video_check_flg = False
if (token == image_token_id) and (video_check_flg is False):
input_token_type.append("image")
elif (token == image_token_id) and (video_check_flg is True):
input_token_type.append("video")
else:
input_token_type.append("text")
input_type_group: list[tuple[str, int, int]] = []
for key, group_iter in itertools.groupby(
enumerate(input_token_type), lambda x: x[1]):
group_list = list(group_iter)
start_index = group_list[0][0]
end_index = group_list[-1][0] + 1
input_type_group.append((key, start_index, end_index))
video_frame_num = 1
mm_data_idx = 0
for modality_type, start_idx, end_idx in input_type_group:
st_idx = llm_pos_ids_list[-1].max() + 1 if len(
llm_pos_ids_list) > 0 else 0
if modality_type == "image":
t, h, w = (
image_grid_thw[mm_data_idx][0],
image_grid_thw[mm_data_idx][1],
image_grid_thw[mm_data_idx][2],
)
llm_grid_t, llm_grid_h, llm_grid_w = \
t, h // spatial_merge_size, w // spatial_merge_size
t_index = torch.arange(llm_grid_t).view(-1, 1).expand(
-1, llm_grid_h * llm_grid_w).flatten()
h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
llm_grid_t, -1, llm_grid_w).flatten()
w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
llm_grid_t, llm_grid_h, -1).flatten()
llm_pos_ids_list.append(
torch.stack([t_index, h_index, w_index]) + st_idx)
mm_data_idx += 1
elif modality_type == "video":
t, h, w = (
video_frame_num,
image_grid_thw[mm_data_idx][1],
image_grid_thw[mm_data_idx][2],
)
llm_grid_t, llm_grid_h, llm_grid_w = \
t, h // spatial_merge_size, w // spatial_merge_size
for t_idx in range(llm_grid_t):
t_index = torch.tensor(t_idx).view(-1, 1).expand(
-1, llm_grid_h * llm_grid_w).flatten()
h_index = torch.arange(llm_grid_h).view(
1, -1, 1).expand(1, -1, llm_grid_w).flatten()
w_index = torch.arange(llm_grid_w).view(
1, 1, -1).expand(1, llm_grid_h, -1).flatten()
llm_pos_ids_list.append(
torch.stack([t_index, h_index, w_index]) + st_idx)
mm_data_idx += 1
video_frame_num += 1
else:
text_len = end_idx - start_idx
llm_pos_ids_list.append(
torch.arange(text_len).view(1, -1).expand(3, -1) +
st_idx)
video_frame_num = 1
else:
text_len = len(input_tokens)
llm_pos_ids_list.append(
torch.arange(text_len).view(1, -1).expand(3, -1))
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
llm_positions = llm_positions[:, context_len:seq_len]
mrope_position_delta = (llm_positions.max() + 1 -
len(input_tokens)).item()
return llm_positions, mrope_position_delta
@classmethod
def _vl_get_input_positions_tensor(
cls,
input_tokens: list[int],
hf_config: PretrainedConfig,
image_grid_thw: Union[list[list[int]], torch.Tensor],
video_grid_thw: Union[list[list[int]], torch.Tensor],
second_per_grid_ts: list[float],
context_len: int = 0,
seq_len: Optional[int] = None,
) -> tuple[torch.Tensor, int]:
"""Get mrope input positions and delta value."""
image_token_id = hf_config.image_token_id
video_token_id = hf_config.video_token_id
vision_start_token_id = hf_config.vision_start_token_id
spatial_merge_size = hf_config.vision_config.spatial_merge_size
tokens_per_second = getattr(hf_config.vision_config,
"tokens_per_second", 1.0)
input_tokens_tensor = torch.tensor(input_tokens)
vision_start_indices = torch.argwhere(
input_tokens_tensor == vision_start_token_id).squeeze(1)
vision_tokens = input_tokens_tensor[vision_start_indices + 1]
image_nums = (vision_tokens == image_token_id).sum()
video_nums = (vision_tokens == video_token_id).sum()
llm_pos_ids_list: list = []
st = 0
remain_images, remain_videos = image_nums, video_nums
image_index, video_index = 0, 0
for _ in range(image_nums + video_nums):
video_second_per_grid_t = 0.0
if image_token_id in input_tokens and remain_images > 0:
ed_image = input_tokens.index(image_token_id, st)
else:
ed_image = len(input_tokens) + 1
if video_token_id in input_tokens and remain_videos > 0:
ed_video = input_tokens.index(video_token_id, st)
else:
ed_video = len(input_tokens) + 1
if ed_image < ed_video:
t, h, w = (
image_grid_thw[image_index][0],
image_grid_thw[image_index][1],
image_grid_thw[image_index][2],
)
image_index += 1
remain_images -= 1
ed = ed_image
else:
t, h, w = (
video_grid_thw[video_index][0],
video_grid_thw[video_index][1],
video_grid_thw[video_index][2],
)
video_second_per_grid_t = 1.0
if second_per_grid_ts:
video_second_per_grid_t = second_per_grid_ts[video_index]
video_index += 1
remain_videos -= 1
ed = ed_video
llm_grid_t, llm_grid_h, llm_grid_w = \
t, h // spatial_merge_size, w // spatial_merge_size
text_len = ed - st
st_idx = llm_pos_ids_list[-1].max() + 1 if len(
llm_pos_ids_list) > 0 else 0
llm_pos_ids_list.append(
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
t_index = (torch.arange(llm_grid_t).view(-1, 1).expand(
-1, llm_grid_h * llm_grid_w) * video_second_per_grid_t *
tokens_per_second).long().flatten()
h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(
llm_grid_t, -1, llm_grid_w).flatten()
w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(
llm_grid_t, llm_grid_h, -1).flatten()
llm_pos_ids_list.append(
torch.stack([t_index, h_index, w_index]) + text_len + st_idx)
st = ed + llm_grid_t * llm_grid_h * llm_grid_w
if st < len(input_tokens):
st_idx = llm_pos_ids_list[-1].max() + 1 if len(
llm_pos_ids_list) > 0 else 0
text_len = len(input_tokens) - st
llm_pos_ids_list.append(
torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
mrope_position_delta = (llm_positions.max() + 1 -
len(input_tokens)).item()
llm_positions = llm_positions[:, context_len:seq_len]
return llm_positions, mrope_position_delta
@classmethod
def _omni_get_input_positions_tensor(
cls,
input_tokens: list[int],
hf_config: PretrainedConfig,
image_grid_thw: Union[list[list[int]], torch.Tensor],
video_grid_thw: Union[list[list[int]], torch.Tensor],
second_per_grid_ts: Optional[list[float]] = None,
context_len: int = 0,
seq_len: Optional[int] = None,
audio_feature_lengths: Optional[torch.Tensor] = None,
use_audio_in_video: bool = False,
) -> tuple[torch.Tensor, int]:
"""Get mrope input positions and delta value (Qwen2.5-Omni version).
Differences from MRotaryEmbedding:
1. Add audio support (and related `audio_feature_lengths`).
2. Add `use_audio_in_video` option to read audio from video inputs.
In this case, audio and vision position ids will be split into
chunks and interleaved.
Example:
(V_i are vision position ids, A_i are audio position ids)
|V_1 ... V_n|A_1 ... A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|...
|vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |...
"""
# TODO(fyabc): refactor and share more code with
# _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
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
vision_end_token_id = thinker_config.vision_end_token_id
seconds_per_chunk = thinker_config.seconds_per_chunk
spatial_merge_size = thinker_config.vision_config.spatial_merge_size
tokens_per_second = getattr(thinker_config.vision_config,
"tokens_per_second", 25)
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)
src_item = input_tokens
audio_seqlens = audio_feature_lengths
if not second_per_grid_ts:
second_per_grid_ts = [1] * video_grid_thw.shape[0]
audio_idx = 0
video_idx = 0
image_idx = 0
new_src_item: list[int] = []
llm_pos_ids_list: list[torch.Tensor] = []
idx = 0
while idx < len(src_item):
new_src_item_len = len(new_src_item)
start_idx = llm_pos_ids_list[-1].max() + 1 if len(
llm_pos_ids_list) > 0 else 0
if src_item[idx] not in [
audio_token_id, video_token_id, image_token_id
]:
if use_audio_in_video and idx > 0:
if src_item[idx] == vision_end_token_id and \
src_item[idx - 1] == audio_end_token_id:
# processing the <|audio_eos|> before <|vision_eos|>
start_idx -= 1
elif src_item[idx] == audio_start_token_id and \
src_item[idx - 1] == vision_start_token_id:
# processing the <|audio_bos|> after <|vision_eos|>
start_idx -= 1
new_src_item.append(src_item[idx])
llm_pos_ids = torch.tensor([start_idx],
dtype=torch.long).expand(3, -1)
llm_pos_ids_list.append(llm_pos_ids)
elif src_item[idx] == audio_token_id:
assert audio_seqlens is not None
audio_seqlen = audio_seqlens[audio_idx]
place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1)
new_src_item.extend([audio_token_id] * place_num)
llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx
llm_pos_ids_list.append(llm_pos_ids)
audio_idx += 1
elif src_item[idx] == image_token_id:
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) * 1 * tokens_per_second).long()
llm_pos_ids = cls._get_llm_pos_ids_for_vision(
start_idx, image_idx, spatial_merge_size, t_index, grid_hs,
grid_ws)
llm_pos_ids_list.append(llm_pos_ids)
vision_seqlen = image_grid_thw[image_idx].prod() // (
spatial_merge_size**2)
new_src_item.extend([image_token_id] * vision_seqlen)
image_idx += 1
elif src_item[idx] == video_token_id and not use_audio_in_video:
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) *
second_per_grid_ts[video_idx] *
tokens_per_second).long()
llm_pos_ids = cls._get_llm_pos_ids_for_vision(
start_idx, video_idx, spatial_merge_size, t_index, grid_hs,
grid_ws)
llm_pos_ids_list.append(llm_pos_ids)
vision_seqlen = video_grid_thw[video_idx].prod() // (
spatial_merge_size**2)
new_src_item.extend([video_token_id] * vision_seqlen)
video_idx += 1
else:
# read audio from video
assert audio_seqlens is not None
audio_seqlen = audio_seqlens[audio_idx]
vision_seqlen = video_grid_thw[video_idx].prod() // (
spatial_merge_size**2)
grid_t = video_grid_thw[video_idx][0]
grid_h = video_grid_thw[video_idx][1]
grid_w = video_grid_thw[video_idx][2]
grid_hs = video_grid_thw[:, 1]
grid_ws = video_grid_thw[:, 2]
t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
t_index = (torch.arange(grid_t) *
second_per_grid_ts[video_idx] *
tokens_per_second).long()
t_index_split_chunk = cls._split_list_into_ranges(
t_index, t_ntoken_per_chunk)
place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2
pure_audio_len = place_num - 2
added_audio_len = 0
audio_llm_pos_ids_list: list[torch.Tensor] = []
for t_chunk in t_index_split_chunk:
vision_ntoken_per_chunk = len(
t_chunk) * grid_h * grid_w // (spatial_merge_size**2)
new_src_item.extend([video_token_id] *
vision_ntoken_per_chunk)
vision_llm_pos_ids_list = cls._get_llm_pos_ids_for_vision(
start_idx, video_idx, spatial_merge_size, t_chunk,
grid_hs, grid_ws).split(1, dim=1)
llm_pos_ids_list.extend(vision_llm_pos_ids_list)
new_src_item.extend(
min(t_ntoken_per_chunk, pure_audio_len -
added_audio_len) * [audio_token_id])
audio_start_idx = start_idx if len(
audio_llm_pos_ids_list
) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1
if min(t_ntoken_per_chunk,
pure_audio_len - added_audio_len) > 0:
audio_llm_pos_ids_list = (torch.arange(
min(t_ntoken_per_chunk, pure_audio_len -
added_audio_len)).expand(3, -1) +
audio_start_idx).split(1,
dim=1)
else:
audio_llm_pos_ids_list = []
added_audio_len += min(t_ntoken_per_chunk,
pure_audio_len - added_audio_len)
llm_pos_ids_list.extend(audio_llm_pos_ids_list)
if added_audio_len < pure_audio_len:
new_src_item.extend(
(pure_audio_len - added_audio_len) * [audio_token_id])
audio_llm_pos_ids_list = (
torch.arange(pure_audio_len - added_audio_len).expand(
3, -1) + llm_pos_ids_list[-1].max() + 1).split(
1, dim=1)
llm_pos_ids_list.extend(audio_llm_pos_ids_list)
audio_idx += 1
video_idx += 1
# move to the next token
idx += len(new_src_item) - new_src_item_len
llm_positions = torch.cat(llm_pos_ids_list, dim=1)
mrope_position_delta = torch.cat(llm_pos_ids_list,
dim=1).max() + 1 - len(src_item)
llm_positions = llm_positions[:, context_len:seq_len]
return llm_positions, mrope_position_delta
@staticmethod
def _get_llm_pos_ids_for_vision(
start_idx: int,
vision_idx: int,
spatial_merge_size: int,
t_index: list[int],
grid_hs: torch.Tensor,
grid_ws: torch.Tensor,
) -> torch.Tensor:
llm_pos_ids_list = []
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())
w_index = (torch.arange(llm_grid_w).view(1, 1, -1).expand(
len(t_index), llm_grid_h, -1).flatten())
t_index_tensor = torch.Tensor(t_index).to(llm_grid_h.device).view(
-1, 1).expand(-1, llm_grid_h * llm_grid_w).long().flatten()
_llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
llm_pos_ids_list.append(_llm_pos_ids + start_idx)
llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
return llm_pos_ids
@staticmethod
def _split_list_into_ranges(lst: torch.Tensor,
interval: int) -> list[list[int]]:
ranges: list[list[int]] = [[]
for _ in range((max(lst) // interval) + 1)]
for num in lst:
index = num // interval
ranges[index].append(num)
return ranges
@staticmethod
def get_next_input_positions(
mrope_position_delta: int,
context_len: int,
seq_len: int,
) -> list[list[int]]:
return [
list(
range(context_len + mrope_position_delta,
seq_len + mrope_position_delta)) for _ in range(3)
]
@staticmethod
def get_next_input_positions_tensor(out: np.ndarray, out_offset: int,
mrope_position_delta: int,
context_len: int, num_new_tokens: int):
values = np.arange(mrope_position_delta + context_len,
mrope_position_delta + context_len + num_new_tokens,
dtype=out.dtype)
out[:, out_offset:out_offset + num_new_tokens] = values
@classmethod
def omni_get_updates_use_audio_in_video(
cls,
thinker_config: PretrainedConfig,
audio_len: int,
video_grid_thw: Union[list[int], torch.Tensor],
video_second_per_grid_t: float,
) -> list[int]:
"""Get video prompt updates when `use_audio_in_video` is True.
In this case, audio and vision update ids will be split into
chunks and interleaved (details in `_omni_get_input_positions_tensor`).
<|video_bos|><|VIDEO|><|video_eos|> =>
<|video_bos|><|audio_bos|>(... chunks ...)<|audio_eos|><|video_eos|>
"""
audio_token_id = thinker_config.audio_token_index
video_token_id = thinker_config.video_token_index
audio_start_token_id = thinker_config.audio_start_token_id
audio_end_token_id = thinker_config.audio_end_token_id
seconds_per_chunk = thinker_config.seconds_per_chunk
spatial_merge_size = thinker_config.vision_config.spatial_merge_size
tokens_per_second = getattr(thinker_config.vision_config,
"tokens_per_second", 25)
grid_t = video_grid_thw[0]
grid_h = video_grid_thw[1]
grid_w = video_grid_thw[2]
t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)
t_index = (torch.arange(grid_t) * video_second_per_grid_t *
tokens_per_second).long()
t_index_split_chunk = cls._split_list_into_ranges(
t_index, t_ntoken_per_chunk)
updates = [audio_start_token_id]
added_audio_len = 0
for t_chunk in t_index_split_chunk:
vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (
spatial_merge_size**2)
updates.extend([video_token_id] * vision_ntoken_per_chunk)
audio_chunk_size = min(t_ntoken_per_chunk,
audio_len - added_audio_len)
updates.extend(audio_chunk_size * [audio_token_id])
added_audio_len += audio_chunk_size
if added_audio_len < audio_len:
updates.extend((audio_len - added_audio_len) * [audio_token_id])
updates.extend([audio_end_token_id])
return updates
vllm_omni/model_executor/models/qwen2_5_omni.py
+710
View File
@@ -0,0 +1,710 @@
from functools import cached_property
from typing import Iterable, List, Optional, Set, Tuple, Union, NamedTuple, Dict
import os
import glob
import numpy as np
import torch
import torch.nn as nn
from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig, Qwen2_5OmniTalkerConfig)
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.layers.linear import ColumnParallelLinear
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.model_executor.models.qwen2_5_omni_thinker import (
Qwen2_5OmniConditionalGenerationMixin,
Qwen2_5OmniThinkerMultiModalProcessor,
Qwen2_5OmniThinkerProcessingInfo,
Qwen2_5OmniThinkerDummyInputsBuilder)
# from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.sequence import IntermediateTensors
from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
init_vllm_registered_model,
maybe_prefix)
from vllm.model_executor.model_loader.weight_utils import download_weights_from_hf
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
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__)
@MULTIMODAL_REGISTRY.register_processor(
Qwen2_5OmniThinkerMultiModalProcessor,
info=Qwen2_5OmniThinkerProcessingInfo,
dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
)
class Qwen2_5OmniForConditionalGeneration(nn.Module, SupportsMultiModal,
SupportsPP,
Qwen2_5OmniConditionalGenerationMixin):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
self.have_multimodal_outputs = True
config: Qwen2_5OmniConfig = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
multimodal_config = vllm_config.model_config.multimodal_config
# keep vllm_config for later submodule init
self.vllm_config = vllm_config
# Initialize thinker components
thinker_config: Qwen2_5OmniThinkerConfig = config.thinker_config
self.thinker_config = thinker_config
self.multimodal_config = multimodal_config
# Initialize talker components
talker_config: Qwen2_5OmniTalkerConfig = config.talker_config
self.talker_config = talker_config
self.model_stage = vllm_config.model_config.model_stage
if self.model_stage=="thinker":
# Initialize thinker model (multimodal processing)
self.thinker = init_vllm_registered_model(
vllm_config=vllm_config,
prefix=maybe_prefix(prefix, "thinker"),
hf_config=thinker_config,
# Use registry architecture key
architectures=["Qwen2_5OmniThinkerModel"],
)
self.model = self.thinker
self.talker = None
self.token2wav = None
elif self.model_stage=="talker":
self.thinker = None
# Initialize talker model wrapper (handles projection + LM)
self.talker = init_vllm_registered_model(
vllm_config=vllm_config,
prefix=maybe_prefix(prefix, "talker"),
hf_config=talker_config,
# Use registry architecture key
architectures=["Qwen2_5OmniTalkerModel"],
)
self.talker.init_multi_modal(thinker_config)
self.model=self.talker
self.token2wav = None
elif self.model_stage=="code2wav":
self.thinker = None
self.talker = None
# Initialize token2wav (code->mel->wav) like thinker/talker
self.token2wav_config = getattr(config, 'token2wav_config', None)
self.token2wav = None
if self.token2wav_config is not None:
self.token2wav = init_vllm_registered_model(
vllm_config=vllm_config,
prefix=maybe_prefix(prefix, "token2wav"),
hf_config=self.token2wav_config,
architectures=["Qwen2_5OmniToken2WavModel"],
)
# voice resources (loaded on demand)
self._token2wav_conds: Dict[str, torch.Tensor] = {}
self._token2wav_ref_mels: Dict[str, torch.Tensor] = {}
self.model = self.token2wav
else:
raise ValueError("Invalid model stage")
# 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:
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,
token2wav_device: Optional[Union[str, torch.device]] = None,
) -> None:
"""Optionally move thinker/talker/token2wav to different devices.
Example:
model.move_submodules_to_devices(
thinker_device='cuda:0',
talker_device='cuda:1',
token2wav_device='cpu',
)
"""
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 token2wav_device is not None and self.token2wav is not None:
self.token2wav.to(token2wav_device)
@cached_property
def sampler(self):
if hasattr(self.model, "sampler"):
return self.model.sampler
return get_sampler()
def get_input_embeddings(
self,
input_ids: torch.Tensor,
multimodal_embeddings=None,
) -> torch.Tensor:
if self.model_stage == "code2wav":
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 thinker model for multimodal processing
return self.model.get_multimodal_embeddings(**kwargs)
def last_index_of(self, list, value):
return len(list) - 1 - list[::-1].index(value)
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 = "Chelsie",
codec: Optional[torch.Tensor] = None,
sampling_metadata: Optional[SamplingMetadata] = None,
logits_index: Optional[int] = None,
sampler = None,
additional_information: Optional[dict[str, object]] = None,
**kwargs: object,
) -> Union[torch.Tensor, IntermediateTensors, OmniOutput]:
"""
Workflow:
1) Thinker: multimodal understanding → text hidden states.
2) If audio requested and codec not provided, use talker to derive codec.
3) If audio requested (or codec provided), use token2wav to synthesize waveform.
4) Return text hidden states (and audio when applicable).
"""
if self.model_stage=="thinker":
# Normalize to batched inputs if caller provides 1D/2D unbatched tensors
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)
#if input_ids is None, set it to a zero tensor, in the length of the same as the embedding seq length
if input_ids is None:
input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev).unsqueeze(0) #(1, 0)
added_batch_dim = True
# 1) Thinker (ensure inputs on thinker's 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
thinker_output = self.thinker(
input_ids=input_ids,
positions=positions[0],
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
**kwargs,
)
if isinstance(thinker_output, tuple):
embeds, text_hidden_states = thinker_output
else:
text_hidden_states = thinker_output
# Text-only path
return OmniOutput(
text_hidden_states=text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states,
multimodal_outputs=None
)
# 2) Talker (if codec not provided)
if self.model_stage=="talker":
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.thinker_reply_part = torch.zeros_like(inputs_embeds)
is_profile = True
else:
is_profile = False
if input_ids is not None and additional_information is not None and not is_profile:
# read from additional_information dict
thinker_result = None
if additional_information is not None and isinstance(additional_information, dict):
thinker_result = additional_information.get("thinker_result")
prompt_embeds = additional_information.get("prompt_embeds")
prompt_token_ids = additional_information.get("prompt_token_ids")
thinker_output_token_ids = additional_information.get("thinker_output_token_ids")
else:
thinker_result = torch.zeros_like(inputs_embeds)
prompt_embeds = torch.zeros_like(inputs_embeds)
prompt_token_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.int64, device=inputs_embeds.device)
thinker_output_token_ids = torch.zeros(inputs_embeds.shape[0], dtype=torch.int64, device=inputs_embeds.device)
if thinker_result is None:
thinker_result = torch.zeros_like(inputs_embeds)
self.thinker_reply_part = thinker_result.squeeze(0)
if self.thinker_reply_part.shape[1] > 1:
self.thinker_reply_part = self.thinker_reply_part[1:, :]
input_ids, inputs_embeds = self._thinker_to_talker_prefill(
voice_type=voice_type,
output_prompt_embeds=thinker_result,
output_token_ids = thinker_output_token_ids,
thinker_prompt_embeds=prompt_embeds,
prompt_token_ids = prompt_token_ids,
)
elif not is_profile:
input_ids, inputs_embeds = \
self._thinker_to_talker_decode_one_step(
output_prompt_embeds=self.thinker_reply_part[:1] if self.thinker_reply_part.shape[0]>=1 \
else torch.zeros(1, self.thinker_reply_part.shape[1]).to(self._module_device(self.model)).to(torch.bfloat16)+(-1.25*2**(-123)),
output_token_ids=input_ids,
)
if self.thinker_reply_part.shape[0] >=1:
self.thinker_reply_part = self.thinker_reply_part[1:, :]
with torch.inference_mode():
talker_hidden = self.talker(
input_ids=input_ids,
positions=positions[0],
inputs_embeds=inputs_embeds,
)
return OmniOutput(
text_hidden_states=talker_hidden,
multimodal_outputs=None
)
if self.model_stage=="code2wav":
code = input_ids if input_ids is not None else torch.zeros(inputs_embeds.shape[0], dtype=torch.long, device=inputs_embeds.device)
audio_tensor = self.generate_audio(code[:-1] if code[-1]==TALKER_CODEC_EOS_TOKEN_ID else code, voice_type)
return OmniOutput(
text_hidden_states = None,
multimodal_outputs = {
"audio": audio_tensor
}
)
return OmniOutput(
text_hidden_states=torch.cat(
[
torch.zeros([inputs_embeds.shape[0],self.talker.config.hidden_size], dtype=torch.bfloat16).to(self._module_device(self.model)),
self.talker.thinker_to_talker_proj(self.talker.get_input_embeddings(torch.tensor([TALKER_CODEC_BOS_TOKEN_ID,TALKER_CODEC_EOS_TOKEN_ID]).to(torch.bfloat16).to(self._module_device(self.model))))[0]
],
dim=0),
multimodal_outputs=None
)
def generate_audio(self, code, voice_type):
token2wav_dev = self._module_device(self.token2wav)
if isinstance(code, torch.Tensor):
code_tensor = code.to(dtype=torch.long, device=token2wav_dev)
else:
code_tensor = torch.as_tensor(code, dtype=torch.long, device=token2wav_dev)
if code_tensor.ndim == 2 and code_tensor.shape[0] == 1:
code_tensor = code_tensor.squeeze(0)
audio_tensor = self._codec_to_audio(code_tensor, voice_type)
return audio_tensor
def _load_talker_embedding(
self,
) -> torch.nn.Embedding:
return self.talker.language_model.model.embed_tokens
def _init_special_tokens_embeddings(
self,
):
# talker embeddings
self.talker_embedding = self._load_talker_embedding()
# embed_text_bos_token
self.tts_text_spk_token_ids = {
# M02: Male voice with standard Mandarin and a slight northern accent
'm02': 151870,
'Ethan': 151870,
# F030: Your anime-styled virtual girlfriend
'f030': 151872,
'Chelsie': 151872,
}
self.default_tts_text_spk_type = list(
self.tts_text_spk_token_ids.keys())[0]
self.tts_text_spk_token_ids['prefix_caching'] = 151870
talker_hf_config = self.talker_config
if hasattr(talker_hf_config, 'talker_config'):
talker_hf_config = talker_hf_config.talker_config
self.embed_text_bos_token = self.thinker_embedding(
torch.tensor(
[talker_hf_config.tts_text_start_token_id],
dtype=torch.long,
device=self._module_device(self.talker),
))
self.embed_text_spk_tokens = {
key:
self.thinker_embedding(
torch.tensor(
[value],
dtype=torch.long,
device=self._module_device(self.talker),
))
for key, value in self.tts_text_spk_token_ids.items()
}
self.embed_text_eos_token = self.thinker_embedding(
torch.tensor(
[talker_hf_config.tts_text_end_token_id],
dtype=torch.long,
device=self._module_device(self.talker),
))
self.embed_text_pad_token = self.thinker_embedding(
torch.tensor(
[talker_hf_config.tts_text_pad_token_id],
dtype=torch.long,
device=self._module_device(self.talker),
))
self.embed_codec_bos_token = self.talker_embedding(
torch.tensor(
[talker_hf_config.tts_codec_start_token_id],
dtype=torch.long,
device=self._module_device(self.talker),
))
self.embed_codec_eos_token = self.talker_embedding(
torch.tensor(
[talker_hf_config.tts_codec_end_token_id],
dtype=torch.long,
device=self._module_device(self.talker),
))
self.embed_codec_pad_token = self.talker_embedding(
torch.tensor(
[talker_hf_config.tts_codec_pad_token_id],
dtype=torch.long,
device=self._module_device(self.talker),
))
return set(["thinker_embedding.weight", "talker_embedding.weight"])
def _get_embed_text_spk_token(self, voice_type: str):
if voice_type not in self.embed_text_spk_tokens:
return self.embed_text_bos_token
return self.embed_text_spk_tokens[voice_type]
def _get_text_spk_token_id(self, voice_type: str):
talker_hf_config = self.talker_config
if hasattr(talker_hf_config, 'talker_config'):
talker_hf_config = talker_hf_config.talker_config
if voice_type not in self.tts_text_spk_token_ids:
return talker_hf_config.tts_text_start_token_id
return self.tts_text_spk_token_ids[voice_type]
def _thinker_to_talker_prefill(
self,
voice_type: str,
output_prompt_embeds,
output_token_ids,
thinker_prompt_embeds,
prompt_token_ids
):
talker_hf_config = self.talker_config
if hasattr(talker_hf_config, 'talker_config'):
talker_hf_config = talker_hf_config.talker_config
# if len(output.outputs[0].token_ids) == 2:
# issue request
prompt_embeds = torch.cat([
thinker_prompt_embeds,
self._get_embed_text_spk_token(voice_type) +
self.embed_codec_pad_token,
output_prompt_embeds[:1] + self.embed_codec_bos_token,
],
dim=0)
prompt_token_ids_processed = prompt_token_ids + [
talker_hf_config.tts_codec_pad_token_id,
output_token_ids[0],
]
input_tokens_len = len(prompt_token_ids_processed)
# the code below is from model runner in Qwen, may need to further discuss later
if input_tokens_len > 2:
prompt_token_ids_processed = (
[self.talker_config.tts_codec_mask_token_id] *
(input_tokens_len - 2) + [
self.talker_config.tts_codec_pad_token_id,
self.talker_config.tts_codec_start_token_id
])
else:
prompt_token_ids_processed = [
self.talker_config.tts_codec_pad_token_id,
self.talker_config.tts_codec_start_token_id,
][-input_tokens_len:]
if isinstance(prompt_token_ids_processed,list):
prompt_token_ids_processed = torch.Tensor(prompt_token_ids_processed).to(torch.int64).to(self._module_device(self.talker))
return prompt_token_ids_processed, prompt_embeds
def _thinker_to_talker_decode_one_step(
self,
output_prompt_embeds,
output_token_ids,
):
processed_output_token_embeds = output_prompt_embeds + self.talker.get_input_embeddings(output_token_ids) #for decode
return output_token_ids, processed_output_token_embeds
def compute_logits(
self,
hidden_states: Union[torch.Tensor, OmniOutput],
sampling_metadata: SamplingMetadata,
) -> Optional[torch.Tensor]:
# Handle OmniOutput type
if isinstance(hidden_states, OmniOutput):
hidden_states = hidden_states.text_hidden_states
# Use thinker model for logits computation
return self.model.compute_logits(hidden_states, sampling_metadata)
def sample(
self,
logits: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> Optional[SamplerOutput]:
# Use thinker model for sampling
return self.model.sample(logits, sampling_metadata)
def generate_speech(self, text_tokens: torch.Tensor, voice_type: str = "default"):
"""
Generate speech from text tokens using the talker and token2wav models.
This method is kept for backward compatibility and direct speech generation.
Args:
text_tokens: Text tokens from thinker model
voice_type: Voice type for speech generation
Returns:
Audio tensor
"""
# Generate codec tokens using talker model
talker_output = self.talker(
input_ids=None,
positions=None,
inputs_embeds=text_tokens
)
# Convert talker output to codec tokens
codec_tokens = self._convert_to_codec_tokens(talker_output)
# Generate audio using token2wav model
return self._codec_to_audio(codec_tokens, voice_type=voice_type)
def _convert_to_codec_tokens(self, talker_output: torch.Tensor, sampling_metadata: SamplingMetadata) -> torch.Tensor:
"""
Reference (HF): use the talker's codec head to obtain logits, suppress BOS,
then greedily select the next codec token for the current step.
"""
with torch.inference_mode():
logits = self.talker.compute_logits(talker_output, None)
if logits is None:
return torch.zeros((talker_output.size(0), 0), dtype=torch.long, device=talker_output.device)
# Suppress only codec_bos, consistent with HF generate's suppress_tokens behavior
bos_id = None
if hasattr(self, 'talker_config') and hasattr(self.talker_config, 'tts_codec_start_token_id'):
bos_id = int(getattr(self.talker_config, 'tts_codec_start_token_id'))
if bos_id is not None:
logits[..., bos_id] = -1e9
# Take the distribution at the last step and select greedily
next_id = self.talker.sample(logits, sampling_metadata).sampled_token_ids
return next_id.to(dtype=torch.long)
def _init_token2wav_model(self):
"""Initialize speaker resources if provided; model is constructed in __init__."""
if self.token2wav is None or self.token2wav_config is None:
return
device = 'cuda' if torch.cuda.is_available() else 'cpu'
# optional speaker resources
conds = getattr(self.token2wav_config, 'conds', None)
ref_mels = getattr(self.token2wav_config, 'ref_mels', None)
if isinstance(conds, dict) and isinstance(ref_mels, dict):
self._token2wav_conds = {k: torch.as_tensor(v, device=device) for k, v in conds.items()}
self._token2wav_ref_mels = {k: torch.as_tensor(v, device=device) for k, v in ref_mels.items()}
# legacy: load from directory if provided
model_path = self.vllm_config.model_config.model
if isinstance(model_path, str) and os.path.isdir(model_path):
spk_pt = os.path.join(model_path, 'spk_dict.pt')
if os.path.exists(spk_pt):
data = torch.load(spk_pt, map_location=device)
for key, value in data.items():
self._token2wav_conds[key] = value["cond"].to(device)
self._token2wav_ref_mels[key] = value["ref_mel"].to(device)
else:
# legacy npy inputs
for f in sorted(glob.glob(os.path.join(model_path, 'inputs', '*spk_emb.npy'))):
key = os.path.basename(f).split('_')[0].lower()
self._token2wav_conds[key] = torch.as_tensor(np.load(f), device=device)
for f in sorted(glob.glob(os.path.join(model_path, 'inputs', '*ref_mel.npy'))):
key = os.path.basename(f).split('_')[0].lower()
self._token2wav_ref_mels[key] = torch.as_tensor(np.load(f), device=device)
def _codec_to_audio(self, codec_tokens: torch.Tensor, voice_type: str = "default") -> Optional[torch.Tensor]:
if self.token2wav is None:
self._init_token2wav_model()
if self.token2wav is None:
return None
# Normalize voice type
voice = (voice_type or 'default')
# Resolve cond / ref_mel if provided
cond = None
ref_mel = None
if voice in self._token2wav_conds and voice in self._token2wav_ref_mels:
cond = self._token2wav_conds[voice]
ref_mel = self._token2wav_ref_mels[voice]
# Fallback: create dummy cond/ref_mel if not provided
token2wav_dev = self._module_device(self.token2wav)
if cond is None:
cond = torch.zeros((1, self.token2wav_config.dit_config.enc_emb_dim), device=token2wav_dev, dtype=torch.float32)
if ref_mel is None:
ref_mel = torch.zeros((1, 300, self.token2wav_config.dit_config.mel_dim), device=token2wav_dev, dtype=torch.float32)
# Ensure codec is (1, T) long tensor on correct device
if isinstance(codec_tokens, torch.Tensor):
codec = codec_tokens.to(dtype=torch.long, device=token2wav_dev)
if codec.ndim == 1:
codec = codec.unsqueeze(0)
else:
codec = torch.as_tensor(codec_tokens, dtype=torch.long, device=token2wav_dev).unsqueeze(0)
# Streaming with chunked process and boundary alignment (rely on token2wav.process_chunk)
factor = getattr(self.token2wav.token2wav.factor, 'factor', 2)
chunk_size = 48
mel_dim = getattr(self.token2wav.token2wav.code2wav_dit_model, 'mel_dim', self.token2wav_config.dit_config.mel_dim)
total_mel = int(codec.shape[1] * factor)
steps = 10
# Prepare initial noise for the whole sequence
y_all = torch.randn((1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev)
logger.info(f"Currently, we do not use the chunked process, we only use the token2wav.process_chunk for the whole sequence.\
The stream mode will be implemented in the future.")
chunk_ends = []
for i in range(codec.shape[1]):
chunk_code_length = i * 2 - 24
finished = i==(codec.shape[1]-1)
if (chunk_code_length > 0 and
chunk_code_length % chunk_size == 0) or finished:
chunk_ends.append(i)
# Number of chunks in mel domain
prev_generated = None
wav_chunks: list = []
prev_id = 0
with torch.inference_mode():
for n,i in enumerate([0]):
finished = (i == codec.shape[1] - 1)
_, audio_chunk = self.token2wav.process_chunk(
conditioning=cond,
reference_mel=ref_mel,
codec_all=codec,
y_all=y_all,
i=n,
steps=steps,
prev_generated=prev_generated if prev_generated is not None else [],
finished=True,
)
prev_generated = audio_chunk
wav_chunks.append(audio_chunk.detach().cpu().numpy())
prev_id = i
if len(wav_chunks) == 0:
return torch.zeros(0, device=token2wav_dev)
waveform = np.concatenate(wav_chunks)
return torch.as_tensor(waveform, device=token2wav_dev)
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 = []
token2wav_weights = []
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('token2wav.'):
token2wav_weights.append((k, v))
else:
raise ValueError(f"Unknown weight prefix: {k}")
# Load thinker weights
if self.thinker:
if thinker_weights:
thinker_loaded = self.thinker.load_weights(thinker_weights)
else:
thinker_loaded = set([k for k,v in thinker_weights])
thinker_loaded = add_prefix_to_loaded_weights(thinker_loaded, 'thinker')
loaded_weights.update(thinker_loaded)
# Load talker weights
if talker_weights and self.talker is not None:
# Map talker weights to appropriate components
if self.thinker is None:
thinker_embedding_weights = [w for n, w in thinker_weights if n == "thinker.model.embed_tokens.weight"]
if thinker_embedding_weights:
self.thinker_embedding = nn.Embedding(thinker_embedding_weights[0].shape[0], thinker_embedding_weights[0].shape[1])
self.thinker_embedding.weight = nn.Parameter(thinker_embedding_weights[0].to(self._module_device(self.talker)))
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 token2wav weights (if any)
if token2wav_weights and self.token2wav is not None:
self._init_token2wav_model()
t2w_loaded = self.token2wav.load_weights(token2wav_weights, os.path.join(self.vllm_config.model_config.model, "spk_dict.pt"))
t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, 'token2wav')
loaded_weights.update(t2w_loaded)
return loaded_weights
vllm_omni/model_executor/models/qwen2_5_omni_talker.py
+239
View File
@@ -0,0 +1,239 @@
from functools import cached_property
from typing import Iterable, List, Optional, Set, Tuple, Union
import torch
import torch.nn as nn
from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniTalkerConfig
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder
from vllm.attention import AttentionMetadata
from vllm.config import VllmConfig
from vllm.forward_context import ForwardContext, get_forward_context
from vllm.model_executor.layers.linear import ColumnParallelLinear
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.model_executor.models.qwen2_5_omni_thinker import (
Qwen2_5OmniConditionalGenerationMixin,
Qwen2_5OmniThinkerMultiModalProcessor,
Qwen2_5OmniThinkerProcessingInfo,
Qwen2_5OmniThinkerDummyInputsBuilder)
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.logger import init_logger
from vllm.sequence import IntermediateTensors
from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP, MultiModalEmbeddings
from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
init_vllm_registered_model,
maybe_prefix,
merge_multimodal_embeddings)
from vllm.model_executor.models.qwen2_5_vl import Qwen2_5_VisionTransformer
@MULTIMODAL_REGISTRY.register_processor(
Qwen2_5OmniThinkerMultiModalProcessor,
info=Qwen2_5OmniThinkerProcessingInfo,
dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
)
class Qwen2_5OmniTalkerForConditionalGeneration(nn.Module, SupportsMultiModal,
SupportsPP,
Qwen2_5OmniConditionalGenerationMixin):
logger = init_logger(__name__)
# Align to thinker-style static mapper for clarity
hf_to_vllm_mapper = WeightsMapper(
orig_to_new_prefix={
# text LM head/body in talker
"talker.codec_head.": "language_model.lm_head.",
"talker.model.": "language_model.model.",
# projection weights
"talker.thinker_to_talker_proj.": "thinker_to_talker_proj.",
# fallback root
"talker.": "",
})
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config: Qwen2_5OmniTalkerConfig = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
self.prefix = prefix
self.quant_config = quant_config
if hasattr(config, "talker_config"):
self.config = config.talker_config
vllm_config.model_config.hf_text_config = \
vllm_config.model_config.hf_config.talker_config
else:
self.config = config
self.thinker_to_talker_proj = ColumnParallelLinear(
self.config.embedding_size,
self.config.hidden_size,
bias=True,
gather_output=True,
skip_bias_add=False,
quant_config=quant_config,
)
self.language_model = init_vllm_registered_model(
vllm_config=vllm_config,
prefix=maybe_prefix(prefix, "language_model"),
hf_config=getattr(self.config, 'text_config', self.config),
architectures=["Qwen2ForCausalLM_old"],
)
self.make_empty_intermediate_tensors = (
self.language_model.make_empty_intermediate_tensors)
def init_multi_modal(self, thinker_config):
self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
self.visual = Qwen2_5_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"),
)
def get_language_model(self) -> torch.nn.Module:
return self.language_model
@cached_property
def sampler(self):
if hasattr(self.language_model, "sampler"):
return self.language_model.sampler
return get_sampler()
def get_input_embeddings(
self,
input_ids: torch.Tensor,
multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
) -> torch.Tensor:
inputs_embeds = self.language_model.get_input_embeddings(input_ids)
if multimodal_embeddings is not None \
and len(multimodal_embeddings) != 0:
# TODO (ywang96): support overlapping modalitiy embeddings so that
# `use_audio_in_video` will work on V1.
inputs_embeds = merge_multimodal_embeddings(
input_ids, inputs_embeds, multimodal_embeddings, [
self.config.image_token_index,
self.config.video_token_index,
self.config.audio_token_index
])
return inputs_embeds
def forward(
self,
input_ids: torch.Tensor = None,
positions: torch.Tensor = None,
intermediate_tensors: Optional[IntermediateTensors] = None,
inputs_embeds: Optional[torch.Tensor] = None,
**kwargs: object,
) -> Union[torch.Tensor, IntermediateTensors]:
assert input_ids is not None or inputs_embeds is not None, "input_ids or inputs_embeds must be provided"
forward_context: ForwardContext = get_forward_context()
attn_metadata: AttentionMetadata = forward_context.attn_metadata
if intermediate_tensors is not None:
inputs_embeds = None
elif inputs_embeds is None:
# for profile_run:
inputs_embeds = self.get_input_embeddings(input_ids)
input_ids = None
# projection
inputs_embeds, _ = self.thinker_to_talker_proj(inputs_embeds)
hidden_states = self.language_model.model(input_ids,
positions,
intermediate_tensors,
inputs_embeds=inputs_embeds)
return hidden_states
def compute_logits(
self,
hidden_states: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> Optional[torch.Tensor]:
return self.language_model.compute_logits(hidden_states,
sampling_metadata)
def sample(
self,
logits: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> Optional[SamplerOutput]:
return self.language_model.sample(logits, sampling_metadata)
def load_weights(self, weights: Iterable[Tuple[str,
torch.Tensor]]) -> Set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=['thinker.', 'token2wav.'],
)
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
self.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
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:
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":
vision_embeddings = self._process_image_input(multimodal_input)
multimodal_embeddings += vision_embeddings
if modality == "video":
video_embeddings = self._process_video_input(multimodal_input)
multimodal_embeddings += video_embeddings
if modality == "audio":
audio_embeddings = self._process_audio_input(multimodal_input)
multimodal_embeddings += audio_embeddings
return multimodal_embeddings
vllm_omni/model_executor/models/qwen2_5_omni_thinker.py
+889
View File
@@ -0,0 +1,889 @@
from collections.abc import Iterable, Mapping, Sequence
from copy import copy
from functools import partial
from typing import Any, Optional, Union
import torch
import torch.nn as nn
from transformers.feature_extraction_utils import BatchFeature
from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig)
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
Qwen2_5OmniAudioEncoder)
from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import (
Qwen2_5OmniProcessor)
from transformers.models.whisper import WhisperFeatureExtractor
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.models.qwen2_5_vl import (
Qwen2_5_VisionTransformer, Qwen2_5_VLImageEmbeddingInputs,
Qwen2_5_VLImageInputs, Qwen2_5_VLImagePixelInputs,
Qwen2_5_VLProcessingInfo, Qwen2_5_VLVideoEmbeddingInputs,
Qwen2_5_VLVideoInputs, Qwen2_5_VLVideoPixelInputs)
from vllm.model_executor.models.qwen2_audio import (
Qwen2AudioInputs, Qwen2AudioProcessingInfo,
_get_feat_extract_output_lengths)
from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.multimodal import MULTIMODAL_REGISTRY
from vllm.multimodal.inputs import (ImageItem, ModalityData,
MultiModalDataDict, MultiModalFieldConfig,
MultiModalKwargs, NestedTensors)
from vllm.multimodal.parse import (AudioProcessorItems, DictEmbeddingItems,
ModalityDataItems, MultiModalDataItems,
MultiModalDataParser)
from vllm.multimodal.processing import (BaseMultiModalProcessor,
PlaceholderFeaturesInfo,
PromptReplacement, PromptUpdate)
from vllm.multimodal.profiling import BaseDummyInputsBuilder
from vllm.sequence import IntermediateTensors
from vllm.transformers_utils.tokenizer import decode_tokens, encode_tokens
from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
from vllm.model_executor.models.utils import (AutoWeightsLoader, WeightsMapper,
init_vllm_registered_model, maybe_prefix,
merge_multimodal_embeddings)
from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
try:
import flash_attn
except (ImportError, ModuleNotFoundError):
flash_attn = None
logger = init_logger(__name__)
def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
audio_feature_lengths = hf_inputs.get("audio_feature_lengths",
torch.empty((0, )))
image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
image_grid_sizes = image_grid_thw.prod(-1)
video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3)))
video_grid_sizes = video_grid_thw.prod(-1)
return dict(
input_audio_features=MultiModalFieldConfig.flat_from_sizes(
"audio", audio_feature_lengths, dim=1),
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
audio_feature_lengths=MultiModalFieldConfig.batched("audio"),
pixel_values=MultiModalFieldConfig.flat_from_sizes(
"image", image_grid_sizes),
image_embeds=MultiModalFieldConfig.flat_from_sizes(
"image", image_grid_sizes),
image_grid_thw=MultiModalFieldConfig.batched("image"),
pixel_values_videos=MultiModalFieldConfig.flat_from_sizes(
"video", video_grid_sizes),
video_embeds=MultiModalFieldConfig.flat_from_sizes(
"video", video_grid_sizes),
video_grid_thw=MultiModalFieldConfig.batched("video"),
second_per_grid_ts=MultiModalFieldConfig.batched("video"),
)
class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
def _parse_audio_data(
self,
data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],
) -> ModalityDataItems[Any, Any]:
if isinstance(data, dict):
return DictEmbeddingItems(
data,
modality="audio",
required_fields={
"input_audio_features", "audio_feature_lengths"
},
fields_factory=_qwen2_5_omni_thinker_field_config,
)
return super()._parse_audio_data(data)
class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo,
Qwen2_5_VLProcessingInfo):
def get_hf_config(self):
return self.ctx.get_hf_config(Qwen2_5OmniConfig).thinker_config
def get_hf_processor(self, **kwargs: object) -> Qwen2_5OmniProcessor:
return self.ctx.get_hf_processor(
Qwen2_5OmniProcessor,
use_fast=kwargs.pop("use_fast", True),
**kwargs,
)
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, Optional[int]]:
return {"audio": None, "image": None, "video": None}
class Qwen2_5OmniThinkerDummyInputsBuilder(
BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]):
def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
num_audios = mm_counts.get("audio", 0)
num_images = mm_counts.get("image", 0)
num_videos = mm_counts.get("video", 0)
hf_processor = self.info.get_hf_processor()
audio_token: str = hf_processor.audio_token
image_token: str = hf_processor.image_token
video_token: str = hf_processor.video_token
return (audio_token * num_audios + image_token * num_images +
video_token * num_videos)
def get_dummy_mm_data(
self,
seq_len: int,
mm_counts: Mapping[str, int],
) -> MultiModalDataDict:
num_audios = mm_counts.get("audio", 0)
num_images = mm_counts.get("image", 0)
num_videos = mm_counts.get("video", 0)
feature_extractor = self.info.get_feature_extractor()
target_audio_length = min(
feature_extractor.chunk_length,
30,
) * feature_extractor.sampling_rate
target_width, target_height = \
self.info.get_image_size_with_most_features()
target_num_frames = \
self.info.get_num_frames_with_most_features(seq_len, mm_counts)
mm_data = {
"audio":
self._get_dummy_audios(length=target_audio_length,
num_audios=num_audios),
"image":
self._get_dummy_images(width=target_width,
height=target_height,
num_images=num_images),
"video":
self._get_dummy_videos(width=target_width,
height=target_height,
num_frames=target_num_frames,
num_videos=num_videos),
}
return mm_data
class Qwen2_5OmniThinkerMultiModalProcessor(
BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]):
def _get_data_parser(self) -> MultiModalDataParser:
feature_extractor = self.info.get_feature_extractor()
return Qwen2_5OmniThinkerMultiModalDataParser(
target_sr=feature_extractor.sampling_rate)
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", [])
# NOTE: WhisperFeatureExtractor cannot handle empty list of audios
if audios:
# NOTE: Qwen2.5-Omni processor accept "audio"
mm_data["audio"] = audios
mm_kwargs = dict(**mm_kwargs, )
hf_inputs = super()._call_hf_processor(
prompt=prompt,
mm_data=mm_data,
mm_kwargs=mm_kwargs,
tok_kwargs=tok_kwargs,
)
input_features = hf_inputs.pop('input_features', None)
feature_attention_mask = hf_inputs.get('feature_attention_mask', None)
if ('input_audio_features' not in hf_inputs
and input_features is not None):
if feature_attention_mask is not None:
input_features = input_features.permute(
0, 2, 1)[feature_attention_mask.bool()].permute(1, 0)
hf_inputs['input_audio_features'] = input_features
if ('audio_feature_lengths' not in hf_inputs
and feature_attention_mask is not None):
hf_inputs['audio_feature_lengths'] = feature_attention_mask.sum(-1)
return hf_inputs
def _get_mm_fields_config(
self,
hf_inputs: BatchFeature,
hf_processor_mm_kwargs: Mapping[str, object],
) -> Mapping[str, MultiModalFieldConfig]:
return _qwen2_5_omni_thinker_field_config(hf_inputs)
def _maybe_apply_prompt_updates(
self,
mm_items: MultiModalDataItems,
hf_processor_mm_kwargs: Mapping[str, object],
prompt_ids: list[int],
mm_kwargs: MultiModalKwargs,
is_update_applied: bool,
) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]:
"""
Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
"""
unbound_prompt_updates = self._get_prompt_updates(
mm_items,
hf_processor_mm_kwargs,
mm_kwargs,
)
mm_prompt_updates = self._bind_and_group_updates(
unbound_prompt_updates)
mm_item_counts = mm_items.get_all_counts()
self._validate_mm_kwargs(mm_kwargs, mm_item_counts)
use_audio_in_video = hf_processor_mm_kwargs.get(
"use_audio_in_video", False)
if is_update_applied:
mm_placeholders = self._find_mm_placeholders(
mm_prompt_updates,
prompt_ids,
mm_item_counts,
)
self._validate_mm_placeholders(
mm_placeholders,
mm_item_counts,
use_audio_in_video=use_audio_in_video)
tokenizer = self.info.get_tokenizer()
prompt = decode_tokens(tokenizer, prompt_ids)
else:
(
prompt_ids,
prompt,
mm_placeholders,
) = self._apply_prompt_updates(
prompt_ids,
mm_prompt_updates,
mm_item_counts,
)
self._validate_mm_placeholders(
mm_placeholders,
mm_item_counts,
use_audio_in_video=use_audio_in_video)
tokenizer = self.info.get_tokenizer()
prompt = decode_tokens(tokenizer, prompt_ids)
if use_audio_in_video:
mm_kwargs["use_audio_in_video"] = True
return prompt_ids, prompt, mm_placeholders
def _get_prompt_updates(
self,
mm_items: MultiModalDataItems,
hf_processor_mm_kwargs: Mapping[str, Any],
out_mm_kwargs: MultiModalKwargs,
) -> 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]
audio_feature_lengths = out_mm_kwargs.get("audio_feature_lengths")
feature_attention_mask = out_mm_kwargs.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
def get_replacement_qwen2_audio(item_idx: int):
item_idx += audio_in_video_item_idx
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_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_in_video_item_idx +
item_idx]
video_grid_thw = out_mm_kwargs["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 = 1.0
return MRotaryEmbedding.omni_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,
)
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 _apply_hf_processor_main(
self,
prompt: Union[str, list[int]],
mm_items: MultiModalDataItems,
hf_processor_mm_kwargs: Mapping[str, object],
tokenization_kwargs: Mapping[str, object],
*,
enable_hf_prompt_update: bool,
) -> tuple[list[int], BatchFeature, bool]:
"""
Qwen2.5-Omni reimplements this function to handle text only.
"""
if isinstance(prompt, str):
if enable_hf_prompt_update:
return self._apply_hf_processor_text_mm(
prompt_text=prompt,
mm_items=mm_items,
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
tokenization_kwargs=tokenization_kwargs,
)
tokenizer = self.info.get_tokenizer()
prompt_ids = encode_tokens(tokenizer, prompt)
else:
prompt_ids = self._apply_hf_processor_tokens_only(prompt)
mm_processed_data = self._apply_hf_processor_mm_only(
mm_items=mm_items,
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
tokenization_kwargs=tokenization_kwargs,
)
return prompt_ids, mm_processed_data, False
def _apply_hf_processor_mm_only(
self,
mm_items: MultiModalDataItems,
hf_processor_mm_kwargs: Mapping[str, object],
tokenization_kwargs: Mapping[str, object],
) -> BatchFeature:
"""
Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`.
"""
mm_counts = mm_items.get_all_counts()
use_audio_in_video = hf_processor_mm_kwargs.get(
"use_audio_in_video", False)
if use_audio_in_video and "video" in mm_counts:
assert "audio" in mm_counts
mm_counts["audio"] -= mm_counts["video"]
_, mm_processed_data, _ = self._apply_hf_processor_text_mm(
prompt_text=self.dummy_inputs.get_dummy_text(mm_counts),
mm_items=mm_items,
hf_processor_mm_kwargs=hf_processor_mm_kwargs,
tokenization_kwargs=tokenization_kwargs,
)
return mm_processed_data
def _validate_mm_placeholders(
self,
mm_placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],
mm_item_counts: Mapping[str, int],
use_audio_in_video: bool = False,
) -> None:
if use_audio_in_video:
mm_item_counts = copy(mm_item_counts)
if "video" in mm_item_counts:
assert "audio" in mm_item_counts
mm_item_counts["audio"] -= mm_item_counts["video"]
super()._validate_mm_placeholders(mm_placeholders, mm_item_counts)
class Qwen2_5OmniConditionalGenerationMixin:
def _validate_and_reshape_mm_tensor(self,
mm_input: object,
name: str,
dim: int = 0) -> torch.Tensor:
if not isinstance(mm_input, (torch.Tensor, list)):
raise ValueError(f"Incorrect type of {name}. "
f"Got type: {type(mm_input)}")
if isinstance(mm_input, torch.Tensor):
return torch.concat(list(mm_input), dim=dim)
else:
return torch.concat(mm_input, dim=dim)
def _parse_and_validate_audio_input(
self, **kwargs: object) -> Optional[Qwen2AudioInputs]:
input_audio_features = kwargs.pop('input_audio_features', None)
audio_feature_lengths = kwargs.pop('audio_feature_lengths', None)
feature_attention_mask = kwargs.pop('feature_attention_mask', None)
if input_audio_features is None:
return None
input_audio_features = self._validate_and_reshape_mm_tensor(
input_audio_features, 'input_audio_features', dim=1)
if feature_attention_mask is not None:
feature_attention_mask = self._validate_and_reshape_mm_tensor(
feature_attention_mask, 'feature_attention_mask')
if not isinstance(input_audio_features, (torch.Tensor, list)):
raise ValueError("Incorrect type of audio input features. "
f"Got type: {type(input_audio_features)}")
return Qwen2AudioInputs(input_features=input_audio_features,
audio_feature_lengths=audio_feature_lengths,
feature_attention_mask=feature_attention_mask)
def _parse_and_validate_image_input(
self,
**kwargs: dict[str, Any],
) -> Optional[Qwen2_5_VLImageInputs]:
pixel_values = kwargs.pop("pixel_values", None)
image_embeds = kwargs.pop("image_embeds", None)
image_grid_thw = kwargs.pop("image_grid_thw", None)
if pixel_values is None and image_embeds is None:
return None
if pixel_values is not None:
pixel_values = self._validate_and_reshape_mm_tensor(
pixel_values, "image pixel values")
image_grid_thw = self._validate_and_reshape_mm_tensor(
image_grid_thw, "image grid_thw")
if not isinstance(pixel_values, (torch.Tensor, list)):
raise ValueError("Incorrect type of image pixel values. "
f"Got type: {type(pixel_values)}")
return Qwen2_5_VLImagePixelInputs(type="pixel_values",
pixel_values=pixel_values,
image_grid_thw=image_grid_thw)
if image_embeds is not None:
image_embeds = self._validate_and_reshape_mm_tensor(
image_embeds, "image embeds")
image_grid_thw = self._validate_and_reshape_mm_tensor(
image_grid_thw, "image grid_thw")
if not isinstance(image_embeds, torch.Tensor):
raise ValueError("Incorrect type of image embeddings. "
f"Got type: {type(image_embeds)}")
return Qwen2_5_VLImageEmbeddingInputs(
type="image_embeds",
image_embeds=image_embeds,
image_grid_thw=image_grid_thw)
def _parse_and_validate_video_input(
self,
**kwargs: dict[str, Any],
) -> Optional[Qwen2_5_VLVideoInputs]:
pixel_values_videos = kwargs.pop("pixel_values_videos", None)
video_embeds = kwargs.pop("video_embeds", None)
video_grid_thw = kwargs.pop("video_grid_thw", None)
if pixel_values_videos is None and video_embeds is None:
return None
if pixel_values_videos is not None:
pixel_values_videos = self._validate_and_reshape_mm_tensor(
pixel_values_videos, "video pixel values")
video_grid_thw = self._validate_and_reshape_mm_tensor(
video_grid_thw, "video grid_thw")
return Qwen2_5_VLVideoPixelInputs(
type="pixel_values_videos",
pixel_values_videos=pixel_values_videos,
video_grid_thw=video_grid_thw,
)
if video_embeds is not None:
video_embeds = self._validate_and_reshape_mm_tensor(
video_embeds, "video embeds")
video_grid_thw = self._validate_and_reshape_mm_tensor(
video_grid_thw, "video grid_thw")
if not isinstance(video_embeds, torch.Tensor):
raise ValueError("Incorrect type of video embeddings. "
f"Got type: {type(video_embeds)}")
return Qwen2_5_VLVideoEmbeddingInputs(
type="video_embeds",
video_embeds=video_embeds,
video_grid_thw=video_grid_thw)
def _process_audio_input(
self,
audio_input: Qwen2AudioInputs,
audio_hashes: list[str] = None,
cached_audio_features: torch.Tensor = None,
) -> torch.Tensor:
input_features = audio_input["input_features"]
audio_feature_lengths = audio_input["audio_feature_lengths"]
if input_features.ndim == 3:
assert input_features.shape[0] == 1
input_features = input_features.squeeze(0)
if audio_feature_lengths.ndim == 2:
assert audio_feature_lengths.shape[
0] == 1 or audio_feature_lengths.shape[1] == 1
if audio_feature_lengths.shape[0] == 1:
audio_feature_lengths = audio_feature_lengths.squeeze(0)
else:
audio_feature_lengths = audio_feature_lengths.squeeze(1)
audio_feat_lengths, audio_output_lengths = (
self.audio_tower._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())
def _process_image_input(
self,
image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]:
if image_input["type"] == "image_embeds":
return image_input["image_embeds"].type(self.visual.dtype)
grid_thw = image_input["image_grid_thw"]
assert grid_thw.ndim == 2
pixel_values = image_input["pixel_values"].type(self.visual.dtype)
image_embeds = self.visual(pixel_values, grid_thw=grid_thw)
# Split concatenated embeddings for each image item.
merge_size = self.visual.spatial_merge_size
sizes = grid_thw.prod(-1) // merge_size // merge_size
return image_embeds.split(sizes.tolist())
def _process_video_input(
self,
video_input: Qwen2_5_VLVideoInputs,
video_hashes: list[str] = None,
cached_video_embeds: torch.Tensor = None) -> torch.Tensor:
if video_input["type"] == "video_embeds":
return video_input["video_embeds"].type(self.visual.dtype)
grid_thw = video_input["video_grid_thw"]
assert grid_thw.ndim == 2
pixel_values_videos = video_input["pixel_values_videos"].type(
self.visual.dtype)
video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw)
# Split concatenated embeddings for each video item.
merge_size = self.visual.spatial_merge_size
sizes = grid_thw.prod(-1) // merge_size // merge_size
return video_embeds.split(sizes.tolist())
@MULTIMODAL_REGISTRY.register_processor(
Qwen2_5OmniThinkerMultiModalProcessor,
info=Qwen2_5OmniThinkerProcessingInfo,
dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder,
)
class Qwen2_5OmniThinkerForConditionalGeneration(
nn.Module, SupportsMultiModal, SupportsPP,
Qwen2_5OmniConditionalGenerationMixin):
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) -> Optional[str]:
if modality.startswith("image"):
return "<|vision_start|><|IMAGE|><|vision_end|>"
if modality.startswith("video"):
return "<|vision_start|><|VIDEO|><|vision_end|>"
if modality.startswith("audio"):
return f"Audio {i}: <|audio_bos|><|AUDIO|><|audio_eos|>"
raise ValueError("Only image, video or audio modality is supported")
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
thinker_config: Qwen2_5OmniThinkerConfig = (
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 = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
self.visual = Qwen2_5_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"),
)
self.quant_config = quant_config
self.language_model = init_vllm_registered_model(
vllm_config=vllm_config,
prefix=maybe_prefix(prefix, "language_model"),
hf_config=thinker_config.text_config,
architectures=["Qwen2ForCausalLM"],
)
self.make_empty_intermediate_tensors = (
self.language_model.make_empty_intermediate_tensors)
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:
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":
vision_embeddings = self._process_image_input(multimodal_input)
multimodal_embeddings += vision_embeddings
if modality == "video":
video_embeddings = self._process_video_input(multimodal_input)
multimodal_embeddings += video_embeddings
if modality == "audio":
audio_embeddings = self._process_audio_input(multimodal_input)
multimodal_embeddings += audio_embeddings
return multimodal_embeddings
def get_input_embeddings(
self,
input_ids: torch.Tensor,
multimodal_embeddings: Optional[MultiModalEmbeddings] = None,
) -> torch.Tensor:
inputs_embeds = self.language_model.get_input_embeddings(input_ids)
if multimodal_embeddings is not None \
and len(multimodal_embeddings) != 0:
# TODO (ywang96): support overlapping modalitiy embeddings so that
# `use_audio_in_video` will work on V1.
inputs_embeds = merge_multimodal_embeddings(
input_ids, inputs_embeds, multimodal_embeddings, [
self.config.image_token_index,
self.config.video_token_index,
self.config.audio_token_index
])
return inputs_embeds
def get_multimodal_embeddings_v0(
self, **kwargs: object) -> Optional[NestedTensors]:
audio_input = self._parse_and_validate_audio_input(**kwargs)
image_input = self._parse_and_validate_image_input(**kwargs)
video_input = self._parse_and_validate_video_input(**kwargs)
if audio_input is None and image_input is None and video_input is None:
return None
multimodal_embeddings: list[tuple[NestedTensors, str]] = []
if audio_input is not None:
audio_embeds = self._process_audio_input(audio_input)
multimodal_embeddings.append((audio_embeds, "audio"))
if image_input is not None:
image_embeds = self._process_image_input(image_input)
multimodal_embeddings.append((image_embeds, "image"))
if video_input is not None:
video_embeds = self._process_video_input(video_input)
multimodal_embeddings.append((video_embeds, "video"))
return multimodal_embeddings
def get_input_embeddings_v0(
self,
input_ids: torch.Tensor,
multimodal_embeddings: Optional[NestedTensors] = None,
) -> torch.Tensor:
inputs_embeds = self.language_model.get_input_embeddings(input_ids)
if multimodal_embeddings is None or len(multimodal_embeddings) == 0:
return inputs_embeds
for embeddings, modality in multimodal_embeddings:
if modality == "audio":
placeholder_token_id = self.config.audio_token_index
if modality == "image":
placeholder_token_id = self.config.image_token_index
if modality == "video":
placeholder_token_id = self.config.video_token_index
inputs_embeds = merge_multimodal_embeddings(
input_ids, inputs_embeds, embeddings, placeholder_token_id)
return inputs_embeds
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
intermediate_tensors: Optional[IntermediateTensors] = None,
inputs_embeds: Optional[torch.Tensor] = None,
**kwargs: object,
) -> Union[torch.Tensor, IntermediateTensors]:
if intermediate_tensors is not None:
inputs_embeds = None
# NOTE: In v1, inputs_embeds is always generated at model runner, this
# condition is for v0 compatibility.
elif inputs_embeds is None:
multimodal_embeddings = self.get_multimodal_embeddings(**kwargs)
inputs_embeds = self.get_input_embeddings(
input_ids, multimodal_embeddings)
text_inputs_embeds = self.get_input_embeddings(input_ids, [(torch.zeros_like(embeddings), modality) for embeddings, modality in multimodal_embeddings] if multimodal_embeddings is not None else None)
input_ids = None
else:
text_inputs_embeds = inputs_embeds
hidden_states = self.language_model.model(input_ids,
positions,
intermediate_tensors,
inputs_embeds=inputs_embeds)
return text_inputs_embeds, hidden_states.unsqueeze(0) #(1, S, D)
def compute_logits(
self,
hidden_states: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> Optional[torch.Tensor]:
return self.language_model.compute_logits(hidden_states,
sampling_metadata)
def load_weights(self, weights: Iterable[tuple[str,
torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=["talker.", "token2wav."],
)
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
vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py
+1690
View File
File diff suppressed because it is too large. Load diff
vllm_omni/model_executor/models/qwen2_old.py
+539
View File
@@ -0,0 +1,539 @@
from typing import Iterable, Optional, Set, Tuple, Union
import torch
from torch import nn
from transformers import Qwen2Config
from vllm.attention import Attention, AttentionType
from vllm.compilation.decorators import support_torch_compile
from vllm.config import CacheConfig, VllmConfig
from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size
from vllm.logger import init_logger
from vllm.model_executor.layers.activation import SiluAndMul
from vllm.model_executor.layers.layernorm import RMSNorm
from vllm.model_executor.layers.linear import (MergedColumnParallelLinear,
QKVParallelLinear,
RowParallelLinear)
from vllm.model_executor.layers.logits_processor import LogitsProcessor
from vllm.model_executor.layers.pooler import Pooler, PoolingType
from vllm.model_executor.layers.quantization import QuantizationConfig
from vllm.model_executor.layers.rotary_embedding import get_rope
from vllm.model_executor.layers.sampler import SamplerOutput, get_sampler
from vllm.model_executor.layers.vocab_parallel_embedding import (
ParallelLMHead, VocabParallelEmbedding)
from vllm.model_executor.model_loader.weight_utils import (
default_weight_loader, maybe_remap_kv_scale_name)
from vllm.v1.pool.metadata import PoolingMetadata
from vllm.model_executor.sampling_metadata import SamplingMetadata
from vllm.sequence import IntermediateTensors, PoolerOutput
from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP
from vllm.model_executor.models.utils import (AutoWeightsLoader, PPMissingLayer, WeightsMapper,
is_pp_missing_parameter,
make_empty_intermediate_tensors_factory, make_layers,
maybe_prefix)
logger = init_logger(__name__)
class Qwen2MLP(nn.Module):
def __init__(
self,
hidden_size: int,
intermediate_size: int,
hidden_act: str,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
self.gate_up_proj = MergedColumnParallelLinear(
hidden_size,
[intermediate_size] * 2,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.gate_up_proj",
)
self.down_proj = RowParallelLinear(
intermediate_size,
hidden_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.down_proj",
)
if hidden_act != "silu":
raise ValueError(f"Unsupported activation: {hidden_act}. "
"Only silu is supported for now.")
self.act_fn = SiluAndMul()
def forward(self, x):
gate_up, _ = self.gate_up_proj(x)
x = self.act_fn(gate_up)
x, _ = self.down_proj(x)
return x
class Qwen2Attention(nn.Module):
def __init__(self,
hidden_size: int,
num_heads: int,
num_kv_heads: int,
max_position: int = 4096 * 32,
rope_theta: float = 10000,
head_dim: Optional[int] = None,
cache_config: Optional[CacheConfig] = None,
quant_config: Optional[QuantizationConfig] = None,
rope_scaling: Optional[Tuple] = None,
prefix: str = "",
attn_type: str = AttentionType.DECODER) -> None:
super().__init__()
self.hidden_size = hidden_size
tp_size = get_tensor_model_parallel_world_size()
self.total_num_heads = num_heads
assert self.total_num_heads % tp_size == 0
self.num_heads = self.total_num_heads // tp_size
self.total_num_kv_heads = num_kv_heads
if self.total_num_kv_heads >= tp_size:
# Number of KV heads is greater than TP size, so we partition
# the KV heads across multiple tensor parallel GPUs.
assert self.total_num_kv_heads % tp_size == 0
else:
# Number of KV heads is less than TP size, so we replicate
# the KV heads across multiple tensor parallel GPUs.
assert tp_size % self.total_num_kv_heads == 0
self.num_kv_heads = max(1, self.total_num_kv_heads // tp_size)
self.head_dim = head_dim or (hidden_size // self.total_num_heads)
self.q_size = self.num_heads * self.head_dim
self.kv_size = self.num_kv_heads * self.head_dim
self.scaling = self.head_dim**-0.5
self.rope_theta = rope_theta
self.qkv_proj = QKVParallelLinear(
hidden_size,
self.head_dim,
self.total_num_heads,
self.total_num_kv_heads,
bias=True,
quant_config=quant_config,
prefix=f"{prefix}.qkv_proj",
)
self.o_proj = RowParallelLinear(
self.total_num_heads * self.head_dim,
hidden_size,
bias=False,
quant_config=quant_config,
prefix=f"{prefix}.o_proj",
)
self.rotary_emb = get_rope(
self.head_dim,
rotary_dim=self.head_dim,
max_position=max_position,
base=self.rope_theta,
rope_scaling=rope_scaling,
)
self.attn = Attention(self.num_heads,
self.head_dim,
self.scaling,
num_kv_heads=self.num_kv_heads,
cache_config=cache_config,
quant_config=quant_config,
prefix=f"{prefix}.attn",
attn_type=attn_type)
def forward(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
) -> torch.Tensor:
qkv, _ = self.qkv_proj(hidden_states)
q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)
q, k = self.rotary_emb(positions, q, k)
attn_output = self.attn(q, k, v)
output, _ = self.o_proj(attn_output)
return output
class Qwen2DecoderLayer(nn.Module):
def __init__(
self,
config: Qwen2Config,
cache_config: Optional[CacheConfig] = None,
quant_config: Optional[QuantizationConfig] = None,
prefix: str = "",
) -> None:
super().__init__()
self.hidden_size = config.hidden_size
# Requires transformers > 4.32.0
rope_theta = getattr(config, "rope_theta", 1000000)
rope_scaling = getattr(config, "rope_scaling", None)
# By default, Qwen2 uses causal attention as it is a decoder-only model.
# You can override the HF config with `is_causal=False` to enable
# bidirectional attention, which is used in some embedding models
# (e.g. Alibaba-NLP/gte-Qwen2-7B-instruct)
if getattr(config, "is_causal", True):
attn_type = AttentionType.DECODER
else:
attn_type = AttentionType.ENCODER_ONLY
self.self_attn = Qwen2Attention(
hidden_size=self.hidden_size,
num_heads=config.num_attention_heads,
max_position=config.max_position_embeddings,
num_kv_heads=config.num_key_value_heads,
rope_theta=rope_theta,
cache_config=cache_config,
quant_config=quant_config,
rope_scaling=rope_scaling,
head_dim=getattr(config, "head_dim", None),
prefix=f"{prefix}.self_attn",
attn_type=attn_type,
)
self.mlp = Qwen2MLP(
hidden_size=self.hidden_size,
intermediate_size=config.intermediate_size,
hidden_act=config.hidden_act,
quant_config=quant_config,
prefix=f"{prefix}.mlp",
)
self.input_layernorm = RMSNorm(config.hidden_size,
eps=config.rms_norm_eps)
self.post_attention_layernorm = RMSNorm(config.hidden_size,
eps=config.rms_norm_eps)
def forward(
self,
positions: torch.Tensor,
hidden_states: torch.Tensor,
residual: Optional[torch.Tensor],
) -> Tuple[torch.Tensor, torch.Tensor]:
# Self Attention
if residual is None:
residual = hidden_states
hidden_states = self.input_layernorm(hidden_states)
else:
hidden_states, residual = self.input_layernorm(
hidden_states, residual)
hidden_states = self.self_attn(
positions=positions,
hidden_states=hidden_states,
)
# Fully Connected
hidden_states, residual = self.post_attention_layernorm(
hidden_states, residual)
hidden_states = self.mlp(hidden_states)
return hidden_states, residual
@support_torch_compile(
dynamic_arg_dims={
"input_ids": 0,
# positions is of shape (3, seq_len) if mrope is enabled for qwen2-vl,
# otherwise (seq_len, ).
"positions": -1,
"intermediate_tensors": 0,
"inputs_embeds": 0,
})
class Qwen2Model(nn.Module):
def __init__(self,
*,
vllm_config: VllmConfig,
prefix: str = "",
decoder_layer_type: type[nn.Module] = Qwen2DecoderLayer):
super().__init__()
config = vllm_config.model_config.hf_config
cache_config = vllm_config.cache_config
quant_config = vllm_config.quant_config
# TODO (@robertgshaw2): see if this can be moved out
if (cache_config.sliding_window is not None
and hasattr(config, "max_window_layers")):
raise ValueError("Sliding window for some but all layers is not "
"supported. This model uses sliding window "
"but `max_window_layers` = {} is less than "
"`num_hidden_layers` = {}. Please open an issue "
"to discuss this feature.".format(
config.max_window_layers,
config.num_hidden_layers,
))
self.config = config
self.quant_config = quant_config
self.vocab_size = config.vocab_size
if get_pp_group().is_first_rank or (config.tie_word_embeddings
and get_pp_group().is_last_rank):
self.embed_tokens = VocabParallelEmbedding(
config.vocab_size,
getattr(config, "embedding_size", config.hidden_size),
quant_config=quant_config,
prefix=f"{prefix}.embed_tokens",
)
else:
self.embed_tokens = PPMissingLayer()
# Use the provided decoder layer type or default to Qwen2DecoderLayer
decoder_layer_type = decoder_layer_type or Qwen2DecoderLayer
self.start_layer, self.end_layer, self.layers = make_layers(
config.num_hidden_layers,
lambda prefix: decoder_layer_type(config=config,
cache_config=cache_config,
quant_config=quant_config,
prefix=prefix),
prefix=f"{prefix}.layers",
)
self.make_empty_intermediate_tensors = (
make_empty_intermediate_tensors_factory(
["hidden_states", "residual"], config.hidden_size))
if get_pp_group().is_last_rank:
self.norm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
else:
self.norm = PPMissingLayer()
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.embed_tokens(input_ids)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
intermediate_tensors: Optional[IntermediateTensors] = None,
inputs_embeds: Optional[torch.Tensor] = None,
) -> Union[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"]
for layer in self.layers[self.start_layer:self.end_layer]:
hidden_states, residual = layer(
positions,
hidden_states,
residual,
)
if not get_pp_group().is_last_rank:
return IntermediateTensors({
"hidden_states": hidden_states,
"residual": residual
})
hidden_states, _ = self.norm(hidden_states, residual)
return hidden_states
def load_weights(self, weights: Iterable[Tuple[str,
torch.Tensor]]) -> Set[str]:
stacked_params_mapping = [
# (param_name, shard_name, shard_id)
("qkv_proj", "q_proj", "q"),
("qkv_proj", "k_proj", "k"),
("qkv_proj", "v_proj", "v"),
("gate_up_proj", "gate_proj", 0),
("gate_up_proj", "up_proj", 1),
]
params_dict = dict(self.named_parameters(remove_duplicate=False))
loaded_params: Set[str] = set()
for name, loaded_weight in weights:
if "rotary_emb.inv_freq" in name:
continue
if (self.quant_config is not None and
(scale_name := self.quant_config.get_cache_scale(name))):
# Loading kv cache quantization scales
param = params_dict[scale_name]
weight_loader = getattr(param, "weight_loader",
default_weight_loader)
loaded_weight = (loaded_weight if loaded_weight.dim() == 0 else
loaded_weight[0])
weight_loader(param, loaded_weight)
loaded_params.add(scale_name)
continue
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)
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
continue
if is_pp_missing_parameter(name, self):
continue
param = params_dict[name]
weight_loader = param.weight_loader
weight_loader(param, loaded_weight, shard_id)
break
else:
# Skip loading extra bias for GPTQ models.
if name.endswith(".bias") and name not in params_dict:
continue
# Remapping the name of FP8 kv-scale.
name = maybe_remap_kv_scale_name(name, params_dict)
if name is None:
continue
if is_pp_missing_parameter(name, self):
continue
param = params_dict[name]
weight_loader = getattr(param, "weight_loader",
default_weight_loader)
weight_loader(param, loaded_weight)
loaded_params.add(name)
return loaded_params
class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
"k_proj",
"v_proj",
],
"gate_up_proj": [
"gate_proj",
"up_proj",
],
}
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
lora_config = vllm_config.lora_config
self.config = config
self.lora_config = lora_config
self.quant_config = quant_config
self.model = Qwen2Model(vllm_config=vllm_config,
prefix=maybe_prefix(prefix, "model"))
if get_pp_group().is_last_rank:
if config.tie_word_embeddings:
self.lm_head = self.model.embed_tokens
else:
self.lm_head = ParallelLMHead(config.vocab_size,
config.hidden_size,
quant_config=quant_config,
prefix=maybe_prefix(
prefix, "lm_head"))
else:
self.lm_head = PPMissingLayer()
self.logits_processor = LogitsProcessor(config.vocab_size)
self.sampler = get_sampler()
self.make_empty_intermediate_tensors = (
self.model.make_empty_intermediate_tensors)
def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:
return self.model.get_input_embeddings(input_ids)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
intermediate_tensors: Optional[IntermediateTensors] = None,
inputs_embeds: Optional[torch.Tensor] = None,
) -> Union[torch.Tensor, IntermediateTensors]:
hidden_states = self.model(input_ids, positions, intermediate_tensors,
inputs_embeds)
return hidden_states
def compute_logits(
self,
hidden_states: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> Optional[torch.Tensor]:
logits = self.logits_processor(self.lm_head, hidden_states,
sampling_metadata)
return logits
def sample(
self,
logits: torch.Tensor,
sampling_metadata: SamplingMetadata,
) -> Optional[SamplerOutput]:
next_tokens = self.sampler(logits, sampling_metadata)
return next_tokens
def load_weights(self, weights: Iterable[Tuple[str,
torch.Tensor]]) -> Set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=(["lm_head."]
if self.config.tie_word_embeddings else None),
)
return loader.load_weights(weights)
class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
packed_modules_mapping = {
"qkv_proj": [
"q_proj",
"k_proj",
"v_proj",
],
"gate_up_proj": [
"gate_proj",
"up_proj",
],
}
hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""})
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
config = vllm_config.model_config.hf_config
quant_config = vllm_config.quant_config
lora_config = vllm_config.lora_config
pooler_config = vllm_config.model_config.pooler_config
self.config = config
self.lora_config = lora_config
self.quant_config = quant_config
self.model = Qwen2Model(vllm_config=vllm_config,
prefix=maybe_prefix(prefix, "model"))
# TODO: Replace this model class with as_embedding_model(
# Qwen2ForCausalLM) after changing the default pooling method
if pooler_config.pooling_type is None:
logger.warning(
"This embedding model will default to last-token pooling in "
"an upcoming version. To avoid breaking changes, you should "
"pass `--override-pooler-config '{\"pooling_type\": \"MEAN\"}'`"
" explicitly.")
self._pooler = Pooler.from_config_with_defaults(
pooler_config,
pooling_type=PoolingType.MEAN,
normalize=True,
softmax=False)
def forward(
self,
input_ids: torch.Tensor,
positions: torch.Tensor,
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> torch.Tensor:
return self.model(input_ids, positions, intermediate_tensors)
def pooler(
self,
hidden_states: torch.Tensor,
pooling_metadata: PoolingMetadata,
) -> Optional[PoolerOutput]:
return self._pooler(hidden_states, pooling_metadata)
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
weights = self.hf_to_vllm_mapper.apply(weights)
weights = ((name, data) for name, data in weights
if not name.startswith("lm_head."))
self.model.load_weights(weights)
vllm_omni/model_executor/models/utils.py
+8
View File
@@ -0,0 +1,8 @@
from vllm.model_executor.models.utils import maybe_prefix
def add_prefix_to_loaded_weights(weights: set[str], prefix: str) -> set[str]:
"""
Add a prefix to the names of the loaded weights.
"""
return {maybe_prefix(prefix, name) for name in weights}
vllm_omni/model_executor/stage_configs/__init__.py
View File
Whitespace-only changes.
vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml
+40
View File
@@ -0,0 +1,40 @@
# stage config for running qwen2.5-omni with architecture of OmniLLM.
stage_args:
- stage_id: 0
engine_args:
model_stage: thinker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
gpu_memory_utilization: 0.32
enforce_eager: true # need to discuss
trust_remote_code: true
engine_output_type: latent # change the param name,such as pooling_output
final_output: true
final_output_type: text
- stage_id: 1
engine_args:
model_stage: talker
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker
scheduler_cls: vllm_omni.core.sched.scheduler.OmniScheduler
gpu_memory_utilization: 0.32
enforce_eager: true
trust_remote_code: true
engine_output_type: latent
engine_input_source: [0]
custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen2_5_omni.thinker2talker
- stage_id: 2
engine_args:
model_stage: code2wav
model_arch: Qwen2_5OmniForConditionalGeneration
worker_cls: vllm_omni.worker.gpu_diffusion_worker.GPUDiffusionWorker
scheduler_cls: vllm_omni.core.sched.diffusion_scheduler.DiffusionScheduler
gpu_memory_utilization: 0.3
enforce_eager: true
trust_remote_code: true
enable_prefix_caching: false
engine_output_type: audio
engine_input_source: [1]
final_output: true
final_output_type: audio
vllm_omni/model_executor/stage_input_processors/__init__.py
View File
Whitespace-only changes.
vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py
+40
View File
@@ -0,0 +1,40 @@
from typing import Union
import torch
from vllm.inputs import TextPrompt
from vllm_omni.inputs.data import OmniTokensPrompt
def thinker2talker(stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None):
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 = []
multi_modal_data = {thinker_output.request_id:
prompt.get('multi_modal_data', None) for thinker_output, prompt in zip(thinker_outputs, prompt)}
for i, thinker_output in enumerate(thinker_outputs):
output = thinker_output.outputs[0]
prompt_token_ids = thinker_output.prompt_token_ids
thinker_output_ids = output.token_ids
prompt_token_ids_len = len(prompt_token_ids)
thinker_hidden_states = output.multimodal_output["latent"].clone().detach().cuda()
talker_inputs.append(
OmniTokensPrompt(
prompt_token_ids=[0] * (len(prompt_token_ids) + 2), # add 2 for codec pad and start token
additional_information={
"thinker_result": thinker_hidden_states[prompt_token_ids_len:].to(torch.float32),
"prompt_embeds": thinker_hidden_states[:prompt_token_ids_len].to(torch.float32),
"prompt_token_ids": prompt_token_ids,
"thinker_output_token_ids": thinker_output_ids,
},
multi_modal_data=multi_modal_data[thinker_output.request_id] if multi_modal_data is not None else None,
mm_processor_kwargs=None,
)
)
return talker_inputs
vllm_omni/worker/gpu_ar_model_runner.py
+578
View File
@@ -0,0 +1,578 @@
"""AR GPU Model Runner for vLLM-omni.
Exposes per-request hidden representations via ModelRunnerOutput.pooler_output
and also outputs sampled tokens.
"""
from __future__ import annotations
from typing import Optional, Union, Any, List
import numpy as np
import torch
from vllm import envs
from vllm.v1.worker.gpu_model_runner import (
EMPTY_MODEL_RUNNER_OUTPUT,
IntermediateTensors,
get_pp_group,
get_tp_group,
has_kv_transfer_group,
set_forward_context,
)
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.multimodal.inputs import MultiModalKwargs
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
from vllm.v1.spec_decode.eagle import EagleProposer
from vllm_omni.engine import PromptEmbedsPayload, AdditionalInformationPayload
from vllm_omni.outputs import OmniModelRunnerOutput
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
class GPUARModelRunner(OmniGPUModelRunner):
"""Autoregressive GPU model runner that returns hidden states per request.
This runner follows the same preparation and forward path as GPUModelRunner
(inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs),
and additionally performs lightweight sampling so that sampled tokens are
available in outputs. Hidden representations are taken at the same indices
that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`).
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.input_ids = torch.zeros(self.max_num_tokens,
dtype=torch.int64,
device=self.device)
@torch.inference_mode()
def execute_model(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
# Update internal state with the new schedule
self._update_states(scheduler_output)
# Decode per-request prompt_embeds / additional_hidden_states 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:
pass
# If there's no work to do, either return empty output or kv-only path
if not scheduler_output.total_num_scheduled_tokens:
if not has_kv_transfer_group():
return EMPTY_MODEL_RUNNER_OUTPUT
return self.kv_connector_no_forward(scheduler_output,
self.vllm_config)
# Prepare decoder inputs and attention metadata
(attn_metadata, attention_cuda_graphs, logits_indices,
spec_decode_metadata, num_scheduled_tokens_np,
spec_decode_common_attn_metadata) = self._prepare_inputs(
scheduler_output)
# Determine number of input tokens for this iteration
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
if (self.use_cuda_graph
and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
num_input_tokens = self.vllm_config.pad_for_cudagraph(
num_scheduled_tokens)
else:
tp_size = self.vllm_config.parallel_config.tensor_parallel_size
if (self.compilation_config.pass_config.enable_sequence_parallelism
and tp_size > 1):
from vllm.utils import round_up # lazy local import
num_input_tokens = round_up(num_scheduled_tokens, tp_size)
else:
num_input_tokens = num_scheduled_tokens
# DP padding
num_pad, num_tokens_across_dp = self.get_dp_padding(num_input_tokens)
num_input_tokens += num_pad
# Multimodal handling (encode and gather embeddings if needed)
if self.is_multimodal_model:
self._execute_mm_encoder(scheduler_output)
mm_embeds = self._gather_mm_embeddings(scheduler_output)
else:
mm_embeds = []
# Always assemble inputs_embeds on first PP rank; overlay per-request prompt_embeds and collect additional_hidden_states for prefill only
if get_pp_group().is_first_rank:
inputs_embeds_scheduled = self.model.get_input_embeddings(
input_ids=self.input_ids[:num_scheduled_tokens],
multimodal_embeddings=(mm_embeds or None)
if self.is_multimodal_model else None,
)
# Copy into persistent buffer to enable CUDA Graph capture
self.inputs_embeds[:num_scheduled_tokens].copy_(
inputs_embeds_scheduled)
# Reset per-step additional information collector
if hasattr(self, "_forward_additional_information"):
self._forward_additional_information = None
# 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)
# For additional_information: handle arbitrary keys
if addi_cpu is not None and isinstance(addi_cpu, dict):
# Lazy init collector dict
if not hasattr(self, "_forward_additional_information") or \
self._forward_additional_information is None:
self._forward_additional_information = {}
# Process tensors (slice by scheduled prompt range) and lists (append per-request)
for k, v in addi_cpu.items():
if isinstance(v, torch.Tensor):
# Slice along token dimension for prefill part
try:
seg = v[num_computed_tokens:
num_computed_tokens + overlay_len].to(
dtype=self.dtype,
device=self.device,
non_blocking=True)
except Exception:
# Fallback: move whole tensor if slicing fails
seg = v.to(dtype=self.dtype,
device=self.device,
non_blocking=True)
prev_val = self._forward_additional_information.get(k)
self._forward_additional_information[k] = (
torch.cat([prev_val, seg], dim=0)
if isinstance(prev_val, torch.Tensor) else seg.clone())
elif isinstance(v, list):
prev_val = self._forward_additional_information.get(k)
if prev_val is None:
self._forward_additional_information[k] = [v]
elif isinstance(prev_val, list):
self._forward_additional_information[k].append(v)
else:
# Mixed types: wrap existing into list
self._forward_additional_information[k] = [prev_val, v]
input_ids = self.input_ids[:num_input_tokens] # preserved for APIs
inputs_embeds = self.inputs_embeds[:num_input_tokens]
model_mm_kwargs = (self._extract_mm_kwargs(scheduler_output)
if self.is_multimodal_model else {})
else:
input_ids = self.input_ids[:num_input_tokens]
inputs_embeds = None
model_mm_kwargs = {}
# Positions/mRoPE
if self.uses_mrope:
positions = self.mrope_positions[:, :num_input_tokens]
else:
positions = self.positions[:num_input_tokens]
# Handle pipeline-parallel intermediate tensors
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)
# Some attention backends only support CUDA Graphs in pure decode.
skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
# Forward pass
with set_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_input_tokens,
num_tokens_across_dp=num_tokens_across_dp,
skip_cuda_graphs=skip_cuda_graphs,
), self.maybe_get_kv_connector_output(
scheduler_output) as kv_connector_output:
model_kwargs_extra = {}
# Only pass additional_information for the prefill part
if hasattr(self, "_forward_additional_information") and \
self._forward_additional_information is not None and \
isinstance(self._forward_additional_information, dict):
model_kwargs_extra["additional_information"] = self._forward_additional_information
model_output = self.model(
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
**MultiModalKwargs.as_kwargs(
model_mm_kwargs,
device=self.device,
),
sampling_metadata=self.input_batch.sampling_metadata,
logits_index=logits_indices,
sampler=self.sampler,
**model_kwargs_extra,
)
if self.use_aux_hidden_state_outputs:
hidden_states, _aux_hidden_states = model_output
else:
hidden_states = model_output
_aux_hidden_states = None
text_hidden_states, multimodal_outputs = (
self.extract_multimodal_outputs(hidden_states))
# Mid PP stages return intermediate tensors unmodified
if not get_pp_group().is_last_rank:
assert isinstance(text_hidden_states, IntermediateTensors)
text_hidden_states.kv_connector_output = kv_connector_output
return text_hidden_states
# Broadcast PP output for external_launcher (torchrun)
broadcast_pp_output = \
self.parallel_config.distributed_executor_backend \
== "external_launcher" and len(get_pp_group().ranks) > 0
if not get_pp_group().is_last_rank:
assert isinstance(text_hidden_states, IntermediateTensors)
if not broadcast_pp_output:
text_hidden_states.kv_connector_output = kv_connector_output
return text_hidden_states
get_pp_group().send_tensor_dict(text_hidden_states.tensors,
all_gather_group=get_tp_group())
logits = None
else:
if self.input_batch.pooling_params:
return self._pool(text_hidden_states, num_scheduled_tokens,
num_scheduled_tokens_np, kv_connector_output)
sample_hidden_states = text_hidden_states[logits_indices]
logits = self.model.compute_logits(sample_hidden_states, None)
if broadcast_pp_output:
model_output_broadcast_data = {
"logits": logits.contiguous(),
} if logits is not None else {}
model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)
assert model_output_broadcast_data is not None
logits = model_output_broadcast_data["logits"]
# Apply structured output bitmasks if present
if scheduler_output.grammar_bitmask is not None:
self.apply_grammar_bitmask(scheduler_output, logits)
# Sample the next token and get logprobs if needed (with spec decode)
sampling_metadata = self.input_batch.sampling_metadata
if spec_decode_metadata is None:
sampler_output = self.sampler(
logits=logits,
sampling_metadata=sampling_metadata,
)
else:
assert logits is not None
bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
sampler_output = self.sampler(
logits=bonus_logits,
sampling_metadata=sampling_metadata,
)
bonus_token_ids = sampler_output.sampled_token_ids
target_logits = logits[spec_decode_metadata.target_logits_indices]
output_token_ids = self.rejection_sampler(
spec_decode_metadata,
None,
target_logits,
bonus_token_ids,
sampling_metadata,
)
sampler_output.sampled_token_ids = output_token_ids
num_nans_in_logits = {}
if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
num_nans_in_logits = self._get_nans_in_logits(logits)
# Handle partial prefill: discard sampled tokens and rewind RNG
discard_sampled_tokens_req_indices = []
for i, req_id in enumerate(self.input_batch.req_ids):
req_state = self.requests[req_id]
seq_len = (req_state.num_computed_tokens +
scheduler_output.num_scheduled_tokens[req_id])
if seq_len < req_state.num_tokens:
generator = self.input_batch.generators.get(i)
if generator is not None:
generator.set_offset(generator.get_offset() - 4)
discard_sampled_tokens_req_indices.append(i)
# Move CPU sync parts
logprobs_tensors = sampler_output.logprobs_tensors
logprobs_lists = logprobs_tensors.tolists() \
if logprobs_tensors is not None else None
# Prompt logprobs if needed
prompt_logprobs_dict = self._get_prompt_logprobs_dict(
text_hidden_states[:num_scheduled_tokens],
scheduler_output,
)
# Parse valid sampled tokens
sampled_token_ids = sampler_output.sampled_token_ids
max_gen_len = sampled_token_ids.shape[-1]
if max_gen_len == 1:
valid_sampled_token_ids = sampled_token_ids.tolist()
else:
valid_sampled_token_ids = self.rejection_sampler.parse_output(
sampled_token_ids,
self.input_batch.vocab_size,
)
for i in discard_sampled_tokens_req_indices:
valid_sampled_token_ids[i].clear()
# Cache sampled tokens
for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
if not sampled_ids:
continue
start_idx = self.input_batch.num_tokens_no_spec[req_idx]
end_idx = start_idx + len(sampled_ids)
self.input_batch.token_ids_cpu[req_idx, start_idx:end_idx] = sampled_ids
self.input_batch.num_tokens_no_spec[req_idx] = end_idx
self.input_batch.num_tokens[req_idx] = end_idx
req_id = self.input_batch.req_ids[req_idx]
req_state = self.requests[req_id]
req_state.output_token_ids.extend(sampled_ids)
# Speculative decoding draft tokens if configured
if not self.speculative_config:
spec_token_ids = None
else:
assert spec_decode_common_attn_metadata is not None
spec_token_ids = self.propose_draft_token_ids(
scheduler_output,
valid_sampled_token_ids,
sampling_metadata,
text_hidden_states,
sample_hidden_states,
_aux_hidden_states,
spec_decode_metadata,
spec_decode_common_attn_metadata,
)
# Convert to per-request tensors on CPU
text_hidden_states_cpu = text_hidden_states.detach().to("cpu").contiguous()
pooler_output: list[Optional[torch.Tensor]] = []
prev_logits_index = 0
for logits_index in logits_indices:
pooler_output.append(text_hidden_states_cpu[prev_logits_index:logits_index+1])
prev_logits_index = logits_index + 1
self.eplb_step()
return OmniModelRunnerOutput(
req_ids=self.input_batch.req_ids,
req_id_to_index=self.input_batch.req_id_to_index,
sampled_token_ids=valid_sampled_token_ids,
spec_token_ids=spec_token_ids,
logprobs=logprobs_lists,
prompt_logprobs_dict=prompt_logprobs_dict,
pooler_output=pooler_output if self.vllm_config.model_config.engine_output_type != "text" else None,
kv_connector_output=kv_connector_output,
num_nans_in_logits=num_nans_in_logits,
)
@torch.inference_mode()
def extract_multimodal_outputs(self, hidden_states: Union[torch.Tensor, List[torch.Tensor]]) -> dict:
if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs:
text_hidden_states = hidden_states.text_hidden_states
multimodal_outputs = hidden_states.multimodal_outputs
elif isinstance(hidden_states, torch.Tensor):
text_hidden_states = hidden_states
multimodal_outputs = {}
elif isinstance(hidden_states, List):
text_hidden_states = hidden_states[0]
multimodal_outputs = {}
else:
raise ValueError(f"Invalid hidden states type: {type(hidden_states)}")
return text_hidden_states, multimodal_outputs
@torch.inference_mode()
def _dummy_run(
self,
num_tokens: int,
capture_attn_cudagraph: bool = False,
skip_eplb: bool = False,
is_profile: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
# Padding for DP
num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
num_tokens += num_pad
# 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)
attn_metadata: Optional[dict[str, Any]] = None
if capture_attn_cudagraph:
attn_metadata = {}
# Make sure max_model_len is used at the graph capture time.
self.seq_lens_np[:num_reqs] = self.max_model_len
self.seq_lens_np[num_reqs:] = 0
self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
non_blocking=True)
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[:num_reqs + 1],
query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs +
1],
seq_lens=self.seq_lens[: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],
num_reqs=num_reqs,
num_actual_tokens=num_tokens,
max_query_len=num_tokens,
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[:num_tokens],
causal=True)
for attn_group in self.attn_groups[kv_cache_group_id]:
attn_metadata_i = attn_group.metadata_builder\
.build_for_cudagraph_capture(common_attn_metadata)
for layer_name in kv_cache_group_spec.layer_names:
attn_metadata[layer_name] = attn_metadata_i
with self.maybe_dummy_run_with_lora(self.lora_config,
num_scheduled_tokens):
if self.is_multimodal_model:
input_ids = None
inputs_embeds = self.inputs_embeds[:num_tokens]
model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)
else:
input_ids = self.input_ids[:num_tokens]
inputs_embeds = None
model_mm_kwargs = {}
if self.uses_mrope:
positions = self.mrope_positions[:, :num_tokens]
else:
positions = self.positions[: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)
with self.maybe_randomize_inputs(input_ids), set_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp):
outputs = self.model(
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
**MultiModalKwargs.as_kwargs(
model_mm_kwargs,
device=self.device,
),
)
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.
if not skip_eplb:
self.eplb_step(is_dummy=True, is_profile=is_profile)
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]
vllm_omni/worker/gpu_ar_worker.py
+67
View File
@@ -0,0 +1,67 @@
from vllm.v1.worker.gpu_worker import Worker as GPUWorker
import torch
import os
import gc
from vllm.utils import GiB_bytes, MemorySnapshot
from vllm.platforms import current_platform
from vllm.worker.worker import _check_if_gpu_supports_dtype
from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
from vllm.model_executor import set_random_seed
from vllm.v1.utils import report_usage_stats
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
class GPUARWorker(GPUWorker):
def init_device(self):
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)
_check_if_gpu_supports_dtype(self.model_config.dtype)
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:
GiB = lambda b: 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}")
# Initialize the distributed environment.
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)
# Construct the model runner
self.model_runner: GPUARModelRunner = GPUARModelRunner(
self.vllm_config, self.device)
if self.rank == 0:
# If usage stat is enabled, collect relevant info.
report_usage_stats(self.vllm_config)
vllm_omni/worker/gpu_diffusion_model_runner.py
+340 -111
View File
@@ -1,133 +1,362 @@
"""Model runner that wraps a diffusers text-to-image pipeline."""
from __future__ import annotations
from dataclasses import dataclass
from typing import Any, Optional
from typing import Optional, Union, List
import numpy as np
import gc
import logging
@dataclass
class DiffusionRunnerOutput:
"""Structured output returned by the diffusion runner."""
import torch
prompt: str
images: list
request_id: str = "diffusion_request"
finished: bool = True
output_type: str = "image"
from vllm.v1.worker.gpu_model_runner import (
EMPTY_MODEL_RUNNER_OUTPUT,
IntermediateTensors,
get_pp_group,
set_forward_context,
)
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
from vllm.multimodal.inputs import MultiModalKwargs
from vllm.v1.core.sched.output import SchedulerOutput
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
from vllm_omni.outputs import OmniModelRunnerOutput
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
class DiffusionModelRunner:
"""Lightweight runner that loads and executes diffusers pipelines."""
logger = logging.getLogger(__name__)
def __init__(
class GPUDiffusionModelRunner(OmniGPUModelRunner):
"""Diffusion 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 diffusion process and returns tensors via `pooler_output`.
"""
@torch.inference_mode()
def execute_model(
self,
model_path: str,
*,
pipeline_name: Optional[str] = None,
device: Optional[str] = None,
dtype: Optional[str] = None,
) -> None:
self.model_path = model_path
self.pipeline_name = pipeline_name or "auto"
self.device, self.torch_dtype = self._resolve_device_and_dtype(device, dtype)
self._pipeline = self._load_pipeline()
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
self._update_states(scheduler_output)
def _resolve_device_and_dtype(
self,
device: Optional[str],
dtype: Optional[str],
):
import torch
if not scheduler_output.total_num_scheduled_tokens:
resolved_device = device
if resolved_device is None:
if hasattr(torch.backends, "mps") and torch.backends.mps.is_available():
resolved_device = "mps"
elif torch.cuda.is_available():
resolved_device = "cuda"
else:
resolved_device = "cpu"
return EMPTY_MODEL_RUNNER_OUTPUT
# Prepare decoder inputs and attention metadata (for batch/order mapping)
(attn_metadata, attention_cuda_graphs, logits_indices,
spec_decode_metadata, num_scheduled_tokens_np,
spec_decode_common_attn_metadata) = self._prepare_inputs(
scheduler_output)
if dtype is not None:
resolved_dtype = getattr(torch, dtype, None)
# 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_across_dp = self.get_dp_padding(num_input_tokens)
num_input_tokens += num_pad
# Multimodal conditioning (e.g., text/audio/video encoders)
if self.is_multimodal_model:
self._execute_mm_encoder(scheduler_output)
mm_embeds = self._gather_mm_embeddings(scheduler_output)
else:
if resolved_device in {"cuda", "mps"}:
resolved_dtype = torch.float16
else:
resolved_dtype = torch.float32
mm_embeds = []
return resolved_device, resolved_dtype
def _load_pipeline(self):
import diffusers
if self.pipeline_name != "auto":
PipelineCls = getattr(diffusers, self.pipeline_name)
# Build inputs to mirror AR runner: input_ids/positions/embeds
if self.is_multimodal_model and get_pp_group().is_first_rank:
inputs_embeds_scheduled = self.model.get_input_embeddings(
input_ids=self.input_ids[:scheduler_output.total_num_scheduled_tokens],
multimodal_embeddings=mm_embeds or None,
)
self.inputs_embeds[:scheduler_output.total_num_scheduled_tokens].copy_(
inputs_embeds_scheduled)
input_ids = self.input_ids[:num_input_tokens]
inputs_embeds = self.inputs_embeds[:num_input_tokens]
model_mm_kwargs = self._extract_mm_kwargs(scheduler_output)
else:
PipelineCls = diffusers.AutoPipelineForText2Image
input_ids = self.input_ids[:num_input_tokens]
inputs_embeds = None
model_mm_kwargs = {}
pipeline = PipelineCls.from_pretrained(
self.model_path,
torch_dtype=self.torch_dtype,
# Positions (mrope or standard)
if self.uses_mrope:
positions = self.mrope_positions[:, :num_input_tokens]
else:
positions = self.positions[:num_input_tokens]
# Intermediate tensors sync for PP (if any)
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)
# Set forward context mainly for resource management and kv connector
skip_cuda_graphs = True # diffusion path does not rely on cuda graphs here
with set_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_input_tokens,
num_tokens_across_dp=num_tokens_across_dp,
skip_cuda_graphs=skip_cuda_graphs,
), self.maybe_get_kv_connector_output(
scheduler_output) as kv_connector_output:
if not get_pp_group().is_last_rank:
# For non-last PP stages, pass through intermediate tensors.
assert intermediate_tensors is not None
intermediate_tensors.kv_connector_output = kv_connector_output
return intermediate_tensors
outputs = self._run_diffusion(
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
multimodal_kwargs=model_mm_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[Optional[torch.Tensor]] = []
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")
self.eplb_step()
return OmniModelRunnerOutput(
req_ids=self.input_batch.req_ids,
req_id_to_index=self.input_batch.req_id_to_index,
sampled_token_ids=[],
spec_token_ids=None,
logprobs=None,
prompt_logprobs_dict={},
pooler_output=pooler_output,
kv_connector_output=kv_connector_output,
num_nans_in_logits={},
)
try:
pipeline = pipeline.to(self.device)
except Exception:
pipeline = pipeline.to("cpu")
self._apply_optimisations(pipeline)
return pipeline
@staticmethod
def _apply_optimisations(pipeline) -> None:
# Enable attention slicing / VAE tiling where available to reduce memory.
try:
if hasattr(pipeline, "enable_attention_slicing"):
pipeline.enable_attention_slicing()
except Exception:
pass
try:
vae = getattr(pipeline, "vae", None)
if vae and hasattr(vae, "enable_tiling"):
vae.enable_tiling()
except Exception:
pass
def generate(
def _run_diffusion(
self,
prompt: str,
*,
height: int = 512,
width: int = 512,
num_inference_steps: int = 30,
guidance_scale: float = 5.0,
negative_prompt: Optional[str] = None,
seed: Optional[int] = None,
image: Optional[Any] = None,
) -> DiffusionRunnerOutput:
import torch
input_ids: torch.Tensor,
positions: torch.Tensor,
intermediate_tensors: Optional[IntermediateTensors],
inputs_embeds: Optional[torch.Tensor],
multimodal_kwargs: dict,
logits_indices: torch.Tensor,
) -> Union[torch.Tensor, list[torch.Tensor]]:
"""Runs the diffusion process and returns per-request tensors.
generator = None
if seed is not None:
try:
generator = torch.Generator(device=self.device).manual_seed(int(seed))
except Exception:
# Some backends (e.g., CPU without full support) may not accept the device.
generator = torch.Generator().manual_seed(int(seed))
output = self._pipeline(
prompt=prompt,
negative_prompt=negative_prompt,
num_inference_steps=int(num_inference_steps),
guidance_scale=float(guidance_scale),
height=int(height),
width=int(width),
generator=generator,
image=image,
Tries model interfaces in the following order for maximal compatibility:
1) model.sample(condition=..., **kwargs)
2) model.forward(condition=..., **kwargs)
3) model.diffuse(condition=..., **kwargs)
"""
# 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,
)
return DiffusionRunnerOutput(
prompt=prompt,
images=getattr(output, "images", []) or [],
)
# For Qwen 2.5 Omni's current implementation, we only support the forward method
if hasattr(self.model, "forward"):
return self.model.forward(**kwargs)
#TODO: add the diffuse method for other models
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_run(
self,
num_tokens: int,
capture_attn_cudagraph: bool = False,
skip_eplb: bool = False,
is_profile: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
# Padding for DP
num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
num_tokens += num_pad
# 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)
attn_metadata: Optional[dict[str, dict]] = None
if capture_attn_cudagraph:
attn_metadata = {}
# Make sure max_model_len is used at the graph capture time.
self.seq_lens_np[:num_reqs] = self.max_model_len
self.seq_lens_np[num_reqs:] = 0
self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
non_blocking=True)
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[:num_reqs + 1],
query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs + 1],
seq_lens=self.seq_lens[: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],
num_reqs=num_reqs,
num_actual_tokens=num_tokens,
max_query_len=num_tokens,
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[:num_tokens],
causal=True)
for attn_group in self.attn_groups[kv_cache_group_id]:
attn_metadata_i = attn_group.metadata_builder \
.build_for_cudagraph_capture(common_attn_metadata)
for layer_name in kv_cache_group_spec.layer_names:
attn_metadata[layer_name] = attn_metadata_i
with self.maybe_dummy_run_with_lora(self.lora_config,
num_scheduled_tokens):
if self.is_multimodal_model:
input_ids = None
inputs_embeds = self.inputs_embeds[:num_tokens]
model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)
else:
input_ids = self.input_ids[:num_tokens]
inputs_embeds = None
model_mm_kwargs = {}
if self.uses_mrope:
positions = self.mrope_positions[:, :num_tokens]
else:
positions = self.positions[: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)
# Diffusion path: avoid CUDA graphs; we only use context for resource wiring
with self.maybe_randomize_inputs(input_ids), set_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp):
outputs = self.model(
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
**MultiModalKwargs.as_kwargs(
model_mm_kwargs,
device=self.device,
),
sampler=None,
)
if self.use_aux_hidden_state_outputs:
hidden_states, _ = outputs
else:
hidden_states = outputs
# Extract multimodal outputs if present; we ignore them here because
# dummy run returns tensors only. The actual diffusion runner returns
# multimodal outputs via pooler_output in execute_model.
text_hidden_states, _ = self.extract_multimodal_outputs(hidden_states)
if not skip_eplb:
self.eplb_step(is_dummy=True, is_profile=is_profile)
logit_indices = np.cumsum(num_scheduled_tokens) - 1
return text_hidden_states, None
@torch.inference_mode()
def _dummy_sampler_run(self, hidden_states: torch.Tensor) -> None:
logger.warning("Dummy sampler run is not implemented for diffusion model")
return None
def profile_run(self) -> None:
# Profile with multimodal encoder & encoder cache, similar to base but
# without any logits/sampler warming.
if self.is_multimodal_model:
mm_budget = self.mm_budget
assert mm_budget is not None
# TODO: handle encoder-decoder models once supported.
if (encoder_budget := mm_budget.get_encoder_budget()) > 0:
(
dummy_modality,
max_tokens,
) = mm_budget.get_modality_with_max_tokens()
(
max_mm_items_per_prompt,
max_mm_items_per_batch,
) = mm_budget.get_max_items(dummy_modality, max_tokens)
batched_dummy_mm_inputs = self._get_mm_dummy_batch(
dummy_modality,
max_mm_items_per_batch,
)
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,
)
self.encoder_cache["tmp"] = dict(
enumerate(dummy_encoder_outputs))
hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True)
if get_pp_group().is_last_rank:
pass # No sampler/pooler warmup for diffusion
self._sync_device()
del hidden_states
self.encoder_cache.clear()
gc.collect()
vllm_omni/worker/gpu_diffusion_worker.py
+61 -50
View File
@@ -1,56 +1,67 @@
"""Worker abstraction for executing diffusion model runners."""
from vllm.v1.worker.gpu_worker import Worker as GPUWorker
import torch
import os
import gc
from vllm.utils import GiB_bytes, MemorySnapshot
from vllm.platforms import current_platform
from vllm.worker.worker import _check_if_gpu_supports_dtype
from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
from vllm.model_executor import set_random_seed
from vllm.v1.utils import report_usage_stats
from __future__ import annotations
from typing import Any, Optional
from vllm.v1.worker.worker_base import WorkerBase
from .gpu_diffusion_model_runner import DiffusionModelRunner, DiffusionRunnerOutput
from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
class DiffusionGPUWorker(WorkerBase):
"""Minimal worker wrapping DiffusionModelRunner.
class GPUDiffusionWorker(GPUWorker):
def init_device(self):
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"
The worker interface mirrors the shape expected by vLLM executors: it owns
a model runner instance and exposes a `generate` method that can be called
by an executor implementation.
"""
# 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)
def __init__(
self,
model_path: str,
*,
pipeline_name: Optional[str] = None,
device: Optional[str] = None,
dtype: Optional[str] = None,
) -> None:
self.model_runner = DiffusionModelRunner(
model_path,
pipeline_name=pipeline_name,
device=device,
dtype=dtype,
)
_check_if_gpu_supports_dtype(self.model_config.dtype)
gc.collect()
torch.cuda.empty_cache()
def generate(
self,
prompt: str,
*,
height: int = 512,
width: int = 512,
num_inference_steps: int = 30,
guidance_scale: float = 5.0,
negative_prompt: Optional[str] = None,
seed: Optional[int] = None,
image: Optional[Any] = None,
) -> DiffusionRunnerOutput:
return self.model_runner.generate(
prompt=prompt,
height=height,
width=width,
num_inference_steps=num_inference_steps,
guidance_scale=guidance_scale,
negative_prompt=negative_prompt,
seed=seed,
image=image,
)
# 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:
GiB = lambda b: 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}")
# Initialize the distributed environment.
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)
# Construct the model runner
self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(
self.vllm_config, self.device)
if self.rank == 0:
# If usage stat is enabled, collect relevant info.
report_usage_stats(self.vllm_config)
vllm_omni/worker/gpu_model_runner.py
+755
View File
@@ -0,0 +1,755 @@
from typing import TYPE_CHECKING, Any, Optional, Union, cast
import numpy as np
import torch
import vllm.envs as envs
from vllm.distributed.parallel_state import get_pp_group, get_tp_group
from vllm.forward_context import set_forward_context
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces_base import VllmModelForPooling
from vllm.multimodal.inputs import MultiModalKwargs
from vllm.sequence import IntermediateTensors
from vllm.utils import LazyLoader
from vllm.v1.attention.backends.utils import CommonAttentionMetadata
from vllm.v1.outputs import EMPTY_MODEL_RUNNER_OUTPUT
from vllm.v1.worker.gpu_input_batch import CachedRequestState
from vllm.sampling_params import SamplingType
from vllm.distributed.kv_transfer import has_kv_transfer_group
from vllm.utils import round_up
from vllm.v1.spec_decode.eagle import EagleProposer
from vllm.v1.worker.gpu_model_runner import GPUModelRunner
from vllm_omni.outputs import OmniModelRunnerOutput
from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
if TYPE_CHECKING:
from vllm.v1.core.sched.output import SchedulerOutput
else:
xgr = LazyLoader("xgr", globals(), "xgrammar")
xgr_torch_compile = LazyLoader(
"xgr_torch_compile", globals(),
"xgrammar.kernels.apply_token_bitmask_inplace_torch_compile")
logger = init_logger(__name__)
class OmniGPUModelRunner(GPUModelRunner):
def _update_states(self, scheduler_output: "SchedulerOutput") -> None:
"""Update the cached states and the persistent batch with the scheduler
output.
The updated states are used by the `_prepare_inputs` function to create
the input GPU tensors for the model.
The SamplingMetadata is updated and copied to the GPU if there is a
new/resumed/paused/finished request in the batch.
"""
# Remove finished requests from the cached states.
for req_id in scheduler_output.finished_req_ids:
self.requests.pop(req_id, None)
self.encoder_cache.pop(req_id, None)
# Remove the finished requests from the persistent batch.
# NOTE(woosuk): There could be an edge case where finished_req_ids and
# scheduled_req_ids overlap. This happens when a request is aborted and
# then resubmitted with the same ID. In this case, we treat them as two
# distinct requests - clearing the cached states for the first request
# and handling the second as a new request.
for req_id in scheduler_output.finished_req_ids:
self.input_batch.remove_request(req_id)
# Free the cached encoder outputs.
for req_id, input_id in scheduler_output.free_encoder_input_ids:
encoder_outputs = self.encoder_cache.get(req_id)
if encoder_outputs is not None:
encoder_outputs.pop(input_id, None)
if not encoder_outputs:
self.encoder_cache.pop(req_id, None)
# Remove the unscheduled requests from the persistent batch.
# NOTE(woosuk): The unscheduled requests are either preempted requests
# or running requests that are not scheduled in this step. We remove
# them from the persistent batch but keep their cached states since
# they will be scheduled again sometime in the future.
scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys()
cached_req_ids = self.input_batch.req_id_to_index.keys()
unscheduled_req_ids = cached_req_ids - scheduled_req_ids
# NOTE(woosuk): The persistent batch optimization assumes that
# consecutive batches contain mostly the same requests. If batches
# have low request overlap (e.g., alternating between two distinct
# sets of requests), this optimization becomes very inefficient.
for req_id in unscheduled_req_ids:
self.input_batch.remove_request(req_id)
req_ids_to_add: list[str] = []
# Add new requests to the cached states.
for new_req_data in scheduler_output.scheduled_new_reqs:
req_id = new_req_data.req_id
sampling_params = new_req_data.sampling_params
pooling_params = new_req_data.pooling_params
if sampling_params and \
sampling_params.sampling_type == SamplingType.RANDOM_SEED:
generator = torch.Generator(device=self.device)
generator.manual_seed(sampling_params.seed)
else:
generator = None
if pooling_params:
assert (task := pooling_params.task) is not None, (
"You did not set `task` in the API")
model = cast(VllmModelForPooling, self.model)
to_update = model.pooler.get_pooling_updates(task)
to_update.apply(pooling_params)
self.requests[req_id] = CachedRequestState(
req_id=req_id,
prompt_token_ids=new_req_data.prompt_token_ids,
mm_inputs=new_req_data.mm_inputs,
mm_positions=new_req_data.mm_positions,
sampling_params=sampling_params,
pooling_params=pooling_params,
generator=generator,
block_ids=new_req_data.block_ids,
num_computed_tokens=new_req_data.num_computed_tokens,
output_token_ids=[],
lora_request=new_req_data.lora_request,
)
# If prompt embeddings are provided, decode and attach to inter_data
try:
if getattr(new_req_data, "prompt_embeds", None) is not None:
payload = new_req_data.prompt_embeds
dtype = getattr(np, payload.dtype)
arr = np.frombuffer(payload.data, dtype=dtype)
arr = arr.reshape(payload.shape)
pe_cpu = torch.from_numpy(arr)
# Store temporarily on CPU; later moved to device in builder
setattr(self.requests[req_id], "prompt_embeds_cpu", pe_cpu)
# Also replace payload with Tensor for user visibility in scheduler_output
try:
new_req_data.prompt_embeds = pe_cpu # type: ignore[assignment]
except Exception:
pass
except Exception:
pass
# Decode additional_information payloads (dictionary)
try:
if getattr(new_req_data, "additional_information", None) is not None:
payload_info = new_req_data.additional_information
info_dict = {}
if isinstance(payload_info, dict):
info_dict = payload_info
else:
from vllm.v1.engine import AdditionalInformationPayload
if 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:
setattr(self.requests[req_id], "additional_information_cpu", info_dict)
except Exception:
pass
# Only relevant for models using M-RoPE (e.g, Qwen2-VL)
if self.uses_mrope:
image_grid_thw = []
video_grid_thw = []
second_per_grid_ts = []
audio_feature_lengths = []
use_audio_in_video = False
for mm_input in self.requests[req_id].mm_inputs:
if mm_input.get("image_grid_thw") is not None:
image_grid_thw.extend(
mm_input["image_grid_thw"].tolist())
if mm_input.get("video_grid_thw") is not None:
video_grid_thw.extend(
mm_input["video_grid_thw"].tolist())
if mm_input.get("second_per_grid_ts") is not None:
second_per_grid_ts.extend(
mm_input["second_per_grid_ts"])
if mm_input.get("audio_feature_lengths") is not None:
audio_feature_lengths.extend(
mm_input["audio_feature_lengths"])
if mm_input.get("use_audio_in_video") is True:
use_audio_in_video = True
hf_config = self.model_config.hf_config
self.requests[req_id].mrope_positions, \
self.requests[req_id].mrope_position_delta = \
MRotaryEmbedding.get_input_positions_tensor(
self.requests[req_id].prompt_token_ids,
hf_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,
)
req_ids_to_add.append(req_id)
# Update the states of the running/resumed requests.
is_last_rank = get_pp_group().is_last_rank
req_data = scheduler_output.scheduled_cached_reqs
for i, req_id in enumerate(req_data.req_ids):
req_state = self.requests[req_id]
num_computed_tokens = req_data.num_computed_tokens[i]
new_block_ids = req_data.new_block_ids[i]
resumed_from_preemption = req_data.resumed_from_preemption[i]
# Update the cached states.
req_state.num_computed_tokens = num_computed_tokens
if not is_last_rank:
# When using PP, the scheduler sends the sampled tokens back,
# because there's no direct communication between the first-
# stage worker and the last-stage worker.
new_token_ids = req_data.new_token_ids[i]
# Add the sampled token(s) from the previous step (if any).
# This doesn't include "unverified" tokens like spec tokens.
num_new_tokens = (num_computed_tokens + len(new_token_ids) -
req_state.num_tokens)
if num_new_tokens == 1:
# Avoid slicing list in most common case.
req_state.output_token_ids.append(new_token_ids[-1])
elif num_new_tokens > 0:
req_state.output_token_ids.extend(
new_token_ids[-num_new_tokens:])
# Update the block IDs.
if not resumed_from_preemption:
# Append the new blocks to the existing block IDs.
for block_ids, new_ids in zip(req_state.block_ids,
new_block_ids):
block_ids.extend(new_ids)
else:
# The request is resumed from preemption.
# Replace the existing block IDs with the new ones.
req_state.block_ids = new_block_ids
req_index = self.input_batch.req_id_to_index.get(req_id)
if req_index is None:
# The request is not in the persistent batch.
# The request was either preempted and resumed later, or was not
# scheduled in the previous step and needs to be added again.
req_ids_to_add.append(req_id)
continue
# Update the persistent batch.
self.input_batch.num_computed_tokens_cpu[req_index] = (
num_computed_tokens)
self.input_batch.block_table.append_row(new_block_ids, req_index)
# For the last rank, we don't need to update the token_ids_cpu
# because the sampled tokens are already cached.
if not is_last_rank:
# Add new_token_ids to token_ids_cpu.
start_token_index = num_computed_tokens
end_token_index = num_computed_tokens + len(new_token_ids)
self.input_batch.token_ids_cpu[
req_index,
start_token_index:end_token_index] = new_token_ids
self.input_batch.num_tokens_no_spec[
req_index] = end_token_index
self.input_batch.num_tokens[req_index] = end_token_index
# Add spec_token_ids to token_ids_cpu.
spec_token_ids = (
scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()))
if spec_token_ids:
num_spec_tokens = len(spec_token_ids)
start_index = self.input_batch.num_tokens_no_spec[req_index]
end_token_index = start_index + num_spec_tokens
self.input_batch.token_ids_cpu[
req_index, start_index:end_token_index] = spec_token_ids
# NOTE(woosuk): `num_tokens` here may include spec tokens.
self.input_batch.num_tokens[req_index] += num_spec_tokens
# Add the new or resumed requests to the persistent batch.
# The smaller empty indices are filled first.
for req_id in req_ids_to_add:
req_state = self.requests[req_id]
self.input_batch.add_request(req_state)
# Condense the batched states if there are gaps left by removed requests
self.input_batch.condense()
# Allow attention backend to reorder the batch, potentially
self._may_reorder_batch(scheduler_output)
# Refresh batch metadata with any pending updates.
self.input_batch.refresh_metadata()
@torch.inference_mode()
def execute_model(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
self._update_states(scheduler_output)
if not scheduler_output.total_num_scheduled_tokens:
if not has_kv_transfer_group():
# Return empty ModelRunnerOutput if there's no work to do.
return EMPTY_MODEL_RUNNER_OUTPUT
return self.kv_connector_no_forward(scheduler_output,
self.vllm_config)
# Prepare the decoder inputs.
(attn_metadata, attention_cuda_graphs, logits_indices,
spec_decode_metadata, num_scheduled_tokens_np,
spec_decode_common_attn_metadata) = (
self._prepare_inputs(scheduler_output))
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
if (self.use_cuda_graph
and num_scheduled_tokens <= self.cudagraph_batch_sizes[-1]):
# Use piecewise CUDA graphs.
# Add padding to the batch size.
num_input_tokens = self.vllm_config.pad_for_cudagraph(
num_scheduled_tokens)
else:
# Eager mode.
# Pad tokens to multiple of tensor_parallel_size when
# enabled collective fusion for SP
tp_size = self.vllm_config.parallel_config.tensor_parallel_size
if self.compilation_config.pass_config. \
enable_sequence_parallelism and tp_size > 1:
num_input_tokens = round_up(num_scheduled_tokens, tp_size)
else:
num_input_tokens = num_scheduled_tokens
# Padding for DP
num_pad, num_tokens_across_dp = 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.is_multimodal_model:
# Run the multimodal encoder if any.
self._execute_mm_encoder(scheduler_output)
mm_embeds = self._gather_mm_embeddings(scheduler_output)
else:
mm_embeds = []
if get_pp_group().is_first_rank:
# Always prepare inputs_embeds on the first PP rank so both
# multimodal and text-only models can consume embeddings.
inputs_embeds_scheduled = self.model.get_input_embeddings(
input_ids=self.input_ids[:num_scheduled_tokens],
multimodal_embeddings=(mm_embeds or None)
if self.is_multimodal_model else None,
)
# TODO(woosuk): Avoid the copy. Optimize.
self.inputs_embeds[:num_scheduled_tokens].copy_(
inputs_embeds_scheduled)
# Overlay per-request custom prompt_embeds for the prompt portion
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)
if pe_cpu is None:
continue
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
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)
input_ids = self.input_ids[:num_input_tokens] # preserved for APIs
inputs_embeds = self.inputs_embeds[:num_input_tokens]
model_mm_kwargs = (self._extract_mm_kwargs(scheduler_output)
if self.is_multimodal_model else {})
else:
# For non-first PP ranks, use token ids as usual; embeddings are
# only consumed on the first rank.
input_ids = self.input_ids[:num_input_tokens]
inputs_embeds = None
model_mm_kwargs = {}
if self.uses_mrope:
positions = self.mrope_positions[:, :num_input_tokens]
else:
positions = self.positions[: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)
# Some attention backends only support CUDA Graphs in pure decode.
# If attention doesn't support CUDA Graphs for this batch, but we
# compiled with full CUDA graphs, we have to skip them entirely.
skip_cuda_graphs = self.full_cuda_graph and not attention_cuda_graphs
# 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,
skip_cuda_graphs=skip_cuda_graphs,
), self.maybe_get_kv_connector_output(
scheduler_output) as kv_connector_output:
model_output = self.model(
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
**MultiModalKwargs.as_kwargs(
model_mm_kwargs,
device=self.device,
),
sampling_metadata=self.input_batch.sampling_metadata,
logits_index=logits_indices,
sampler=self.sampler,
)
if self.use_aux_hidden_state_outputs:
hidden_states, aux_hidden_states = model_output
else:
hidden_states = model_output
aux_hidden_states = None
text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
# Broadcast PP output for external_launcher (torchrun)
# to make sure we are synced across pp ranks
# TODO: Support overlapping mirco-batches
# https://github.com/vllm-project/vllm/issues/18019
broadcast_pp_output = \
self.parallel_config.distributed_executor_backend \
== "external_launcher" and len(get_pp_group().ranks) > 0
if not get_pp_group().is_last_rank:
# For mid-pipeline stages, return the hidden states.
assert isinstance(text_hidden_states, IntermediateTensors)
if not broadcast_pp_output:
text_hidden_states.kv_connector_output = kv_connector_output
return text_hidden_states
get_pp_group().send_tensor_dict(text_hidden_states.tensors,
all_gather_group=get_tp_group())
logits = None
else:
if self.input_batch.pooling_params:
return self._pool(text_hidden_states, num_scheduled_tokens,
num_scheduled_tokens_np, kv_connector_output)
sample_hidden_states = text_hidden_states[logits_indices]
logits = self.model.compute_logits(sample_hidden_states, None)
if broadcast_pp_output:
model_output_broadcast_data = {
"logits": logits.contiguous(),
} if logits is not None else {}
model_output_broadcast_data = get_pp_group().broadcast_tensor_dict(
model_output_broadcast_data, src=len(get_pp_group().ranks) - 1)
assert model_output_broadcast_data is not None
logits = model_output_broadcast_data["logits"]
# Apply structured output bitmasks if present
if scheduler_output.grammar_bitmask is not None:
self.apply_grammar_bitmask(scheduler_output, logits)
# Sample the next token and get logprobs if needed.
sampling_metadata = self.input_batch.sampling_metadata
if spec_decode_metadata is None:
sampler_output = self.sampler(
logits=logits,
sampling_metadata=sampling_metadata,
)
else:
# When indexing with a tensor (bonus_logits_indices), PyTorch
# creates a new tensor with separate storage from the original
# logits tensor. This means any in-place operations on bonus_logits
# won't affect the original logits tensor.
assert logits is not None
bonus_logits = logits[spec_decode_metadata.bonus_logits_indices]
sampler_output = self.sampler(
logits=bonus_logits,
sampling_metadata=sampling_metadata,
)
bonus_token_ids = sampler_output.sampled_token_ids
# Just like `bonus_logits`, `target_logits` is a new tensor with
# separate storage from the original `logits` tensor. Therefore,
# it is safe to update `target_logits` in place.
target_logits = logits[spec_decode_metadata.target_logits_indices]
output_token_ids = self.rejection_sampler(
spec_decode_metadata,
None, # draft_probs
target_logits,
bonus_token_ids,
sampling_metadata,
)
sampler_output.sampled_token_ids = output_token_ids
num_nans_in_logits = {}
if envs.VLLM_COMPUTE_NANS_IN_LOGITS:
num_nans_in_logits = self._get_nans_in_logits(logits)
# TODO(woosuk): The following loop can be slow since it iterates over
# the requests one by one. Optimize.
discard_sampled_tokens_req_indices = []
for i, req_id in enumerate(self.input_batch.req_ids):
req_state = self.requests[req_id]
seq_len = (req_state.num_computed_tokens +
scheduler_output.num_scheduled_tokens[req_id])
if seq_len < req_state.num_tokens:
# Ignore the sampled token for partial prefills.
# Rewind the generator state as if the token was not sampled.
# This relies on cuda-specific torch-internal impl details
generator = self.input_batch.generators.get(i)
if generator is not None:
generator.set_offset(generator.get_offset() - 4)
# Record the index of the request that should not be sampled,
# so that we could clear the sampled tokens before returning.
discard_sampled_tokens_req_indices.append(i)
# NOTE: GPU -> CPU Sync happens here.
# Move as many CPU operations as possible before this sync point.
logprobs_tensors = sampler_output.logprobs_tensors
logprobs_lists = logprobs_tensors.tolists() \
if logprobs_tensors is not None else None
# Compute prompt logprobs if needed.
prompt_logprobs_dict = self._get_prompt_logprobs_dict(
text_hidden_states[:num_scheduled_tokens],
scheduler_output,
)
# Get the valid generated tokens.
import os
sampled_token_ids = sampler_output.sampled_token_ids
max_gen_len = sampled_token_ids.shape[-1]
if max_gen_len == 1:
# No spec decode tokens.
valid_sampled_token_ids = sampled_token_ids.tolist()
else:
# Includes spec decode tokens.
valid_sampled_token_ids = self.rejection_sampler.parse_output(
sampled_token_ids,
self.input_batch.vocab_size,
)
# Mask out the sampled tokens that should not be sampled.
for i in discard_sampled_tokens_req_indices:
valid_sampled_token_ids[i].clear()
# Cache the sampled tokens in the model runner, so that the scheduler
# doesn't need to send them back.
# NOTE(woosuk): As an exception, when using PP, the scheduler sends
# the sampled tokens back, because there's no direct communication
# between the first-stage worker and the last-stage worker.
for req_idx, sampled_ids in enumerate(valid_sampled_token_ids):
if not sampled_ids:
continue
start_idx = self.input_batch.num_tokens_no_spec[req_idx]
end_idx = start_idx + len(sampled_ids)
assert end_idx <= self.max_model_len, (
"Sampled token IDs exceed the max model length. "
f"Total number of tokens: {end_idx} > max_model_len: "
f"{self.max_model_len}")
self.input_batch.token_ids_cpu[req_idx,
start_idx:end_idx] = sampled_ids
self.input_batch.num_tokens_no_spec[req_idx] = end_idx
self.input_batch.num_tokens[req_idx] = end_idx
req_id = self.input_batch.req_ids[req_idx]
req_state = self.requests[req_id]
req_state.output_token_ids.extend(sampled_ids)
if not self.speculative_config:
# Speculative decoding is not enabled.
spec_token_ids = None
else:
assert spec_decode_common_attn_metadata is not None
spec_token_ids = self.propose_draft_token_ids(
scheduler_output,
valid_sampled_token_ids,
sampling_metadata,
text_hidden_states,
sample_hidden_states,
aux_hidden_states,
spec_decode_metadata,
spec_decode_common_attn_metadata,
)
self.eplb_step()
return OmniModelRunnerOutput(
req_ids=self.input_batch.req_ids,
req_id_to_index=self.input_batch.req_id_to_index,
sampled_token_ids=valid_sampled_token_ids,
spec_token_ids=spec_token_ids,
logprobs=logprobs_lists,
prompt_logprobs_dict=prompt_logprobs_dict,
pooler_output=[],
kv_connector_output=kv_connector_output,
num_nans_in_logits=num_nans_in_logits,
multimodal_outputs=multimodal_outputs,
)
@torch.inference_mode()
def extract_multimodal_outputs(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, dict]:
if hasattr(self.model, "have_multimodal_outputs") and self.model.have_multimodal_outputs:
text_hidden_states = hidden_states.text_hidden_states
multimodal_outputs = hidden_states.multimodal_outputs
else:
text_hidden_states = hidden_states
multimodal_outputs = {}
return text_hidden_states, multimodal_outputs
@torch.inference_mode()
def _dummy_run(
self,
num_tokens: int,
capture_attn_cudagraph: bool = False,
skip_eplb: bool = False,
is_profile: bool = False,
) -> tuple[torch.Tensor, torch.Tensor]:
# Padding for DP
num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens)
num_tokens += num_pad
# 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)
attn_metadata: Optional[dict[str, Any]] = None
if capture_attn_cudagraph:
attn_metadata = {}
# Make sure max_model_len is used at the graph capture time.
self.seq_lens_np[:num_reqs] = self.max_model_len
self.seq_lens_np[num_reqs:] = 0
self.seq_lens[:num_reqs].copy_(self.seq_lens_cpu[:num_reqs],
non_blocking=True)
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[:num_reqs + 1],
query_start_loc_cpu=self.query_start_loc_cpu[:num_reqs +
1],
seq_lens=self.seq_lens[: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],
num_reqs=num_reqs,
num_actual_tokens=num_tokens,
max_query_len=num_tokens,
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[:num_tokens],
causal=True)
for attn_group in self.attn_groups[kv_cache_group_id]:
attn_metadata_i = attn_group.metadata_builder\
.build_for_cudagraph_capture(common_attn_metadata)
for layer_name in kv_cache_group_spec.layer_names:
attn_metadata[layer_name] = attn_metadata_i
with self.maybe_dummy_run_with_lora(self.lora_config,
num_scheduled_tokens):
if self.is_multimodal_model:
input_ids = None
inputs_embeds = self.inputs_embeds[:num_tokens]
model_mm_kwargs = self._dummy_mm_kwargs(num_reqs)
else:
input_ids = self.input_ids[:num_tokens]
inputs_embeds = None
model_mm_kwargs = {}
if self.uses_mrope:
positions = self.mrope_positions[:, :num_tokens]
else:
positions = self.positions[: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)
with self.maybe_randomize_inputs(input_ids), set_forward_context(
attn_metadata,
self.vllm_config,
num_tokens=num_tokens,
num_tokens_across_dp=num_tokens_across_dp):
outputs = self.model(
input_ids=input_ids,
positions=positions,
intermediate_tensors=intermediate_tensors,
inputs_embeds=inputs_embeds,
**MultiModalKwargs.as_kwargs(
model_mm_kwargs,
device=self.device,
),
sampler = None
)
if self.use_aux_hidden_state_outputs:
hidden_states, _ = outputs
else:
hidden_states = outputs
text_hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
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.
if not skip_eplb:
self.eplb_step(is_dummy=True, is_profile=is_profile)
logit_indices = np.cumsum(num_scheduled_tokens) - 1
return text_hidden_states, text_hidden_states[logit_indices]
vllm_omni/worker/model_runner.py
+199
View File
@@ -0,0 +1,199 @@
from dataclasses import dataclass
from typing import Optional, Dict, Any, List
import itertools
import torch
from vllm.worker.model_runner import ModelInputForGPU, ModelInputForGPUBuilder
from vllm.utils import flatten_2d_lists, async_tensor_h2d
from vllm.lora.layers import LoRAMapping
from vllm.sequence import MultiModalKwargs
@dataclass(frozen=True)
class OmniModelInputForGPU(ModelInputForGPU):
# New: optional additional information dict for current scheduled tokens
additional_information: Optional[Dict[str, Any]] = None
class OmniModelInputForGPUBuilder(ModelInputForGPUBuilder):
def build(self) -> OmniModelInputForGPU:
"""Finalize the builder intermediate data and
create on-device tensors.
"""
# Combine and flatten intermediate data.
input_tokens = list[int]()
inputs_embeds_list = list[torch.Tensor]()
token_types = list[int]()
for inter_data in self.inter_data_list:
for cur_input_tokens in inter_data.input_tokens:
input_tokens.extend(cur_input_tokens)
for cur_token_types in inter_data.token_types:
token_types.extend(cur_token_types)
if inter_data.inputs_embeds is not None:
inputs_embeds_list.append(
inter_data.inputs_embeds.to(
dtype=self.runner.model_config.dtype,
device=self.runner.device))
# Support v1 direct-transfer prompt embeds attached on request state
if inter_data.inputs_embeds is None:
try:
# Locate req_id -> request state to fetch decoded CPU embeds
req_id = inter_data.request_id # type: ignore[attr-defined]
req_state = getattr(self.runner, "requests", {}).get(req_id)
pe_cpu = getattr(req_state, "prompt_embeds_cpu", None)
if pe_cpu is not None:
inter_data.inputs_embeds = pe_cpu.to(
dtype=self.runner.model_config.dtype,
device=self.runner.device)
inputs_embeds_list.append(inter_data.inputs_embeds)
except Exception:
pass
inputs_embeds: Optional[torch.Tensor]
if len(inputs_embeds_list) == 0:
inputs_embeds = None
else:
inputs_embeds = torch.cat(inputs_embeds_list, dim=0).to(
dtype=self.runner.model_config.dtype,
device=self.runner.device)
assert len(inputs_embeds) == len(input_tokens)
if not input_tokens and inputs_embeds is None:
# This may happen when all prefill requests hit
# prefix caching and there is no decode request.
return self.model_input_cls()
mrope_input_positions: Optional[List[List[int]]] = None
if any(inter_data.mrope_input_positions is not None
for inter_data in self.inter_data_list):
mrope_input_positions = [[] for _ in range(3)]
for idx in range(3):
for inter_data in self.inter_data_list:
msections = inter_data.mrope_input_positions
if msections is None:
for _seq_input_positions in inter_data.input_positions:
mrope_input_positions[idx].extend(
_seq_input_positions)
else:
for _seq_mrope_input_positions in msections:
mrope_input_positions[idx].extend(
_seq_mrope_input_positions[idx])
input_positions = None
else:
input_positions = []
for inter_data in self.inter_data_list:
for cur_input_positions in inter_data.input_positions:
input_positions.extend(cur_input_positions)
seq_lens = []
query_lens = []
max_decode_seq_len = 0
max_encoder_seq_len = 0
for inter_data in self.inter_data_list:
seq_lens.extend(inter_data.seq_lens)
query_lens.extend(inter_data.query_lens)
if not inter_data.is_prompt:
max_decode_seq_len = max(max_decode_seq_len,
max(inter_data.seq_lens))
if self.runner.model_config.is_encoder_decoder:
max_encoder_seq_len = max(max_encoder_seq_len,
inter_data.encoder_seq_len)
# Mapping from request IDs to sequence IDs. Used for Jamba models
# that manages the cache by itself.
request_ids_to_seq_ids = {
data.request_id: data.seq_ids
for data in self.inter_data_list
}
cuda_graph_pad_size = self._get_cuda_graph_pad_size(
num_seqs=len(seq_lens),
max_decode_seq_len=max_decode_seq_len,
max_encoder_seq_len=max_encoder_seq_len)
batch_size = len(input_tokens)
if cuda_graph_pad_size != -1:
# If cuda graph can be used, pad tensors accordingly.
# See `capture_model` API for more details.
# vLLM uses cuda graph only for decoding requests.
batch_size += cuda_graph_pad_size
# Tokens and positions.
if cuda_graph_pad_size:
input_tokens.extend(itertools.repeat(0, cuda_graph_pad_size))
assert self.runner.device is not None
input_tokens_tensor = async_tensor_h2d(input_tokens, torch.long,
self.runner.device,
self.runner.pin_memory)
token_types_tensor = async_tensor_h2d(token_types, torch.long,
self.runner.device,
self.runner.pin_memory) \
if token_types else None
if mrope_input_positions is not None:
for idx in range(3):
mrope_input_positions[idx].extend(
itertools.repeat(0, cuda_graph_pad_size))
input_positions_tensor = async_tensor_h2d(mrope_input_positions,
torch.long,
self.runner.device,
self.runner.pin_memory)
else:
input_positions.extend(itertools.repeat(0, cuda_graph_pad_size))
input_positions_tensor = async_tensor_h2d(input_positions,
torch.long,
self.runner.device,
self.runner.pin_memory)
# Sequence and query lengths.
if cuda_graph_pad_size:
seq_lens.extend(itertools.repeat(1, cuda_graph_pad_size))
# Attention metadata.
attn_metadata = self.attn_metadata_builder.build(
seq_lens, query_lens, cuda_graph_pad_size, batch_size)
# LoRA data.
lora_requests = set()
lora_mapping = None
if self.enable_lora:
lora_requests = set(r for data in self.inter_data_list
for r in data.lora_requests)
lora_index_mapping = flatten_2d_lists([
flatten_2d_lists(inter_data.lora_index_mapping)
for inter_data in self.inter_data_list
])
if cuda_graph_pad_size:
lora_index_mapping.extend(
itertools.repeat(0, cuda_graph_pad_size))
lora_prompt_mapping = flatten_2d_lists([
flatten_2d_lists(inter_data.lora_prompt_mapping)
for inter_data in self.inter_data_list
])
lora_mapping = LoRAMapping(
**dict(index_mapping=lora_index_mapping,
prompt_mapping=lora_prompt_mapping,
is_prefill=not self.decode_only))
# Multi-modal data.
multi_modal_kwargs_list = [
data.multi_modal_kwargs for data in self.inter_data_list
if data.multi_modal_kwargs is not None
]
multi_modal_kwargs = MultiModalKwargs.batch(multi_modal_kwargs_list)
return self.model_input_cls(
input_tokens=input_tokens_tensor,
inputs_embeds=inputs_embeds,
input_positions=input_positions_tensor,
token_types=token_types_tensor,
attn_metadata=attn_metadata,
seq_lens=seq_lens,
query_lens=query_lens,
lora_mapping=lora_mapping,
lora_requests=lora_requests,
multi_modal_kwargs=multi_modal_kwargs,
request_ids_to_seq_ids=request_ids_to_seq_ids,
finished_requests_ids=self.finished_requests_ids)