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

[CI] fix ci and format existing code (#71)

Browse Source 
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
This commit is contained in:
Jiangyun Zhu authored and GitHub committed 2025-11-18 20:54:26 +08:00
1 parent 09fab2ecc9
commit 5707fc78d5
52 files changed
+813 -2155
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

.github/workflows/ci-all-files.yml
+2 -2
View File
@@ -39,8 +39,8 @@ jobs:
run: |
python -m pip install --upgrade pip uv
uv venv --python "${{ matrix.python-version }}" --seed
echo "VIRTUAL_ENV=$PWD/.venv" >> $GITHUB_ENV
echo "PATH=$PWD/.venv/bin:$PATH" >> $GITHUB_ENV
echo "VIRTUAL_ENV=${PWD}/.venv" >> "$GITHUB_ENV"
echo "PATH=${PWD}/.venv/bin:$PATH" >> "$GITHUB_ENV"
source .venv/bin/activate
uv pip install -e ".[dev]"
.github/workflows/ci.yml
+5 -5
View File
@@ -68,8 +68,8 @@ jobs:
run: |
python -m pip install --upgrade pip uv
uv venv --python "${{ matrix.python-version }}" --seed
echo "VIRTUAL_ENV=$PWD/.venv" >> $GITHUB_ENV
echo "PATH=$PWD/.venv/bin:$PATH" >> $GITHUB_ENV
echo "VIRTUAL_ENV=${PWD}/.venv" >> "$GITHUB_ENV"
echo "PATH=${PWD}/.venv/bin:${PATH}" >> "$GITHUB_ENV"
# Activate env to ensure subsequent commands succeed in this step
source .venv/bin/activate
uv pip install -e ".[dev]"
@@ -86,8 +86,8 @@ jobs:
TEST_CHANGED_ONLY="false"
fi
echo "test_changed_only=$TEST_CHANGED_ONLY" >> $GITHUB_OUTPUT
echo "enable_detection=$ENABLE_CHANGED_FILES_DETECTION" >> $GITHUB_OUTPUT
echo "test_changed_only=$TEST_CHANGED_ONLY" >> "$GITHUB_OUTPUT"
echo "enable_detection=$ENABLE_CHANGED_FILES_DETECTION" >> "$GITHUB_OUTPUT"
echo "Test scope: $TEST_CHANGED_ONLY"
echo "Changed files detection: $ENABLE_CHANGED_FILES_DETECTION"
@@ -107,7 +107,7 @@ jobs:
- name: Run tests (changed files only)
if: steps.test-scope.outputs.test_changed_only == 'true' && steps.changed-tests.outputs.has_tests == 'true'
run: |
pytest ${{ steps.changed-tests.outputs.test_files }} -v -m "not slow"
pytest "${{ steps.changed-tests.outputs.test_files }}" -v -m "not slow"
- name: Run tests (all files)
if: steps.test-scope.outputs.test_changed_only == 'false' || steps.changed-tests.outputs.has_tests == 'false'
.github/workflows/pre-commit.yml
-2
View File
@@ -22,5 +22,3 @@ jobs:
python-version: "3.12"
cache: 'pip'
- uses: pre-commit/action@v3.0.1
with:
extra_args: --hook-stage manual
.pre-commit-config.yaml
+4 -13
View File
@@ -12,21 +12,12 @@ repos:
- id: trailing-whitespace
args: ["--markdown-linebreak-ext=md"]
- repo: https://github.com/PyCQA/isort
rev: 5.13.2
hooks:
- id: isort
- repo: https://github.com/psf/black
rev: 24.10.0
hooks:
- id: black
- repo: https://github.com/astral-sh/ruff-pre-commit
rev: v0.11.7
rev: v0.14.0
hooks:
- id: ruff
args: [--output-format, github, --fix, --line-length=160]
- id: ruff-check
args: [--output-format, github, --fix]
- id: ruff-format
- repo: https://github.com/crate-ci/typos
rev: v1.35.5
README.md
+1 -1
View File
@@ -85,7 +85,7 @@ uv pip install --no-build-isolation --editable .
```
### Verification for successful installation of vLLM
Just run the command below. If no error, it demonstrates that the installation is successfull.
Just run the command below. If no error, it demonstrates that the installation is successful.
```bash
python -c "import vllm._C"
```
docs/contributing/design_documents/vllm_omni_design.md
+3 -3
View File
@@ -206,7 +206,7 @@ dit_stage_config.executor_class = DiffusersPipelineExecutor
Similar to OmniLLM in offline inference, add some asynchronous processing, referring to AsyncLLM
```python
from vllm.v1.engine.async_llm improt AsyncLLM
from vllm.v1.engine.async_llm import AsyncLLM
class AsyncOmniLLM(AsyncLLM):
@@ -317,13 +317,13 @@ class DiffusersPipelineExecutor(Executor):
# Generate image
prompt_embeds = self._get_and_process_prompt_embeds(scheduler_output, positive_magic)
negtive_prompt_embeds = self.pipe.embed_prompt(" ")
negative_prompt_embeds = self.pipe.embed_prompt(" ")
width, height = aspect_ratios["16:9"]
image = pipe(
prompt_embeds=prompt_embeds,
negtive_prompt_embeds=negtive_prompt_embeds,
negative_prompt_embeds=negative_prompt_embeds,
width=width,
height=height,
num_inference_steps=50,
docs/mkdocs/hooks/generate_examples.py
+2 -8
View File
@@ -115,9 +115,7 @@ class Example:
".cfg",
".ini",
}
is_other_file = lambda file: ( # noqa: E731
file.is_file() and file != self.main_file and file.suffix in text_extensions
)
is_other_file = lambda file: (file.is_file() and file != self.main_file and file.suffix in text_extensions) # noqa: E731
return [file for file in self.path.rglob("*") if is_other_file(file)]
def determine_title(self) -> str:
@@ -175,11 +173,7 @@ class Example:
if self.is_code:
# Use relative path from ROOT_DIR for snippet inclusion
main_file_rel_path = self.main_file.relative_to(ROOT_DIR)
content += (
f"{code_fence}{self.main_file.suffix[1:]}\n"
f'--8<-- "{main_file_rel_path}"\n'
f"{code_fence}\n"
)
content += f'{code_fence}{self.main_file.suffix[1:]}\n--8<-- "{main_file_rel_path}"\n{code_fence}\n'
else:
with open(self.main_file, encoding="utf-8") as f:
# Skip the title from md snippets as it's been included above
docs/user_guide/examples/offline_inference/qwen_2_5_omni.md
+1 -1
View File
@@ -52,7 +52,7 @@ uv pip install --no-build-isolation --editable .
```
### Verification for successful installation of vLLM
Just run the command below. If no error, it demonstrates that the installation is successfull.
Just run the command below. If no error, it demonstrates that the installation is successful.
```bash
python -c "import vllm._C"
```
examples/offline_inference/qwen2_5_omni/end2end.py
+11 -29
View File
@@ -7,8 +7,8 @@ import numpy as np
import soundfile as sf
import torch
from utils import make_omni_prompt
from vllm.sampling_params import SamplingParams
from vllm_omni.entrypoints.omni_llm import OmniLLM
_os_env_toggle.environ["VLLM_USE_V1"] = "1"
@@ -45,29 +45,17 @@ def parse_args():
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", nargs="+", default=None, help="Input text prompts."
)
parser.add_argument(
"--voice-type", default="default", help="Voice type, e.g., m02, f030, default."
)
parser.add_argument("--hf-revision", default=None, help="Optional HF revision (branch/tag/commit).")
parser.add_argument("--prompts", nargs="+", default=None, 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("--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(
"--init-sleep-seconds",
@@ -163,7 +151,7 @@ def main():
try:
# Preferred: load from txt file (one prompt per line)
if getattr(args, "txt_prompts", None) and args.prompt_type == "text":
with open(args.txt_prompts, "r", encoding="utf-8") as f:
with open(args.txt_prompts, encoding="utf-8") as f:
lines = [ln.strip() for ln in f.readlines()]
args.prompts = [ln for ln in lines if ln != ""]
print(f"[Info] Loaded {len(args.prompts)} prompts from {args.txt_prompts}")
@@ -172,9 +160,7 @@ def main():
raise
if args.prompts is None:
raise ValueError(
"No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)"
)
raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts <file.txt> (with --prompt_type text)")
omni_llm = OmniLLM(
model=model_name,
log_stats=args.enable_stats,
@@ -223,9 +209,7 @@ def main():
omni_outputs = omni_llm.generate(prompt, sampling_params_list)
# Determine output directory: prefer --output-dir; fallback to --output-wav
output_dir = (
args.output_dir if getattr(args, "output_dir", None) else args.output_wav
)
output_dir = args.output_dir if getattr(args, "output_dir", None) else args.output_wav
os.makedirs(output_dir, exist_ok=True)
for stage_outputs in omni_outputs:
if stage_outputs.final_output_type == "text":
@@ -251,9 +235,7 @@ def main():
request_id = int(output.request_id)
audio_tensor = output.multimodal_output["audio"]
output_wav = os.path.join(output_dir, f"output_{output.request_id}.wav")
sf.write(
output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000
)
sf.write(output_wav, audio_tensor.detach().cpu().numpy(), samplerate=24000)
print(f"Request ID: {request_id}, Saved audio to {output_wav}")
examples/offline_inference/qwen2_5_omni/extract_prompts.py
+1 -1
View File
@@ -28,7 +28,7 @@ def main() -> None:
args = parser.parse_args()
prompts = []
with open(args.input, "r", encoding="utf-8", errors="ignore") as f:
with open(args.input, encoding="utf-8", errors="ignore") as f:
for line in f:
if len(prompts) >= args.topk:
break
examples/offline_inference/qwen2_5_omni/processing_omni.py
+11 -35
View File
@@ -72,8 +72,7 @@ def smart_resize(
"""
if max(height, width) / min(height, width) > MAX_RATIO:
raise ValueError(
f"absolute aspect ratio must be smaller than {MAX_RATIO}, "
f"got {max(height, width) / min(height, width)}"
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))
@@ -88,9 +87,7 @@ def smart_resize(
return h_bar, w_bar
def fetch_image(
ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR
) -> Image.Image:
def fetch_image(ele: dict[str, str | Image.Image], size_factor: int = IMAGE_FACTOR) -> Image.Image:
if "image" in ele:
image = ele["image"]
else:
@@ -110,10 +107,7 @@ def fetch_image(
else:
image_obj = Image.open(image)
if image_obj is None:
raise ValueError(
f"Unrecognized image input, support local path, http url, "
f"base64 and PIL.Image, got {image}"
)
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:
@@ -163,9 +157,7 @@ def smart_nframes(
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`"
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:
@@ -179,10 +171,7 @@ def smart_nframes(
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}], "
f"but got {nframes}."
)
raise ValueError(f"nframes should in interval [{FRAME_FACTOR}, {total_frames}], but got {nframes}.")
return nframes
@@ -204,10 +193,7 @@ def _read_video_torchvision(
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."
)
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()
@@ -221,8 +207,7 @@ def _read_video_torchvision(
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=}, "
f"duration={total_duration}s, time={time.time() - st:.3f}s"
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()
@@ -258,15 +243,10 @@ def _read_video_decord(
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."
)
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=}, "
f"time={time.time() - st:.3f}s"
)
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()
@@ -295,9 +275,7 @@ def get_video_reader_backend() -> str:
return video_reader_backend
def fetch_video(
ele: dict, image_factor: int = IMAGE_FACTOR
) -> torch.Tensor | list[Image.Image]:
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)
@@ -341,9 +319,7 @@ def fetch_video(
process_info.pop("type", None)
process_info.pop("video", None)
images = [
fetch_image(
{"image": video_element, **process_info}, size_factor=image_factor
)
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)
examples/offline_inference/qwen2_5_omni/utils.py
+13 -26
View File
@@ -1,5 +1,5 @@
import tempfile
from typing import Dict, List, Optional, Union
from typing import Optional, Union
from urllib.request import urlopen
import librosa
@@ -10,17 +10,16 @@ import torch
import torchvision.io
from processing_omni import fetch_image, fetch_video
from transformers import AutoConfig, AutoProcessor
from vllm.inputs import TextPrompt
from vllm_omni.inputs.data import OmniTokensPrompt
# Simple caches to avoid repeated heavy HF loads per prompt
_PROCESSOR_CACHE: Dict[str, "AutoProcessor"] = {}
_CONFIG_CACHE: Dict[str, "AutoConfig"] = {}
_PROCESSOR_CACHE: dict[str, "AutoProcessor"] = {}
_CONFIG_CACHE: dict[str, "AutoConfig"] = {}
def get_system_prompt():
return {
"role": "system",
"content": [
@@ -47,7 +46,6 @@ def resample_wav_to_16khz(input_filepath):
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,
@@ -73,10 +71,8 @@ def fetch_and_read_video(args, video_url: str, fps=2):
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}
)
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
@@ -96,8 +92,7 @@ def fetch_and_read_video(args, video_url: str, fps=2):
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}, "
f"status_code:{resp.status_code}, resp:{resp}"
f"Failed to fetch video from {video_url}, status_code:{resp.status_code}, resp:{resp}"
)
temp_video_file.write(urlopen(video_url).read())
@@ -111,11 +106,10 @@ def fetch_and_read_video(args, video_url: str, fps=2):
def make_inputs_qwen2_omni(
args,
messages: List[Dict[str, Union[str, List[Dict[str, str]]]]],
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
# Cached processor/config to prevent per-prompt reloading and repeated warnings
@@ -134,7 +128,7 @@ def make_inputs_qwen2_omni(
# Decide legacy flag only once based on config (default True if unknown)
if getattr(args, "legacy_omni_video", None) is None:
if config is not None and hasattr(config, "architectures"):
args.legacy_omni_video = not ("Qwen2_5OmniModel" in config.architectures)
args.legacy_omni_video = "Qwen2_5OmniModel" not in config.architectures
else:
args.legacy_omni_video = True
@@ -181,9 +175,7 @@ def make_inputs_qwen2_omni(
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"
):
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"
@@ -199,7 +191,7 @@ def make_inputs_qwen2_omni(
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))
raise ValueError(f"Unknown ele {ele}")
# insert a audio after the video
message["content"].insert(
index + 1,
@@ -296,10 +288,7 @@ def make_audio_in_video_v2_prompt(args):
"content": [
{
"type": "video_url",
"video_url": (
"https://qianwen-res.oss-cn-beijing.aliyuncs.com/"
"Qwen2.5-Omni/draw_small.mp4"
),
"video_url": ("https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen2.5-Omni/draw_small.mp4"),
},
],
},
@@ -313,9 +302,7 @@ def make_audio_in_video_v2_prompt(args):
return prompt
def make_omni_prompt(
args, prompt=None
) -> Union[OmniTokensPrompt, List[OmniTokensPrompt]]:
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":
examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py
-1
View File
@@ -16,7 +16,6 @@ SEED = 42
def get_system_prompt():
return {
"role": "system",
"content": [
pyproject.toml
+13 -26
View File
@@ -77,33 +77,8 @@ vllm-omni = "vllm_omni.entrypoints.cli.main:main"
where = ["."]
include = ["vllm_omni*"]
[tool.black]
line-length = 160
target-version = ['py312', 'py313']
include = '\.pyi?$'
extend-exclude = '''
/(
# directories
\.eggs
| \.git
| \.hg
| \.mypy_cache
| \.tox
| \.venv
| build
| dist
)/
'''
[tool.isort]
profile = "black"
multi_line_output = 3
line_length = 160
known_first_party = ["vllm_omni"]
[tool.ruff]
line-length = 160
target-version = "py312, py313"
line-length = 120
exclude = [
".eggs",
".git",
@@ -128,6 +103,10 @@ select = [
ignore = [
"E203", # whitespace before ':' (conflicts with black)
# W503 is not needed in ruff as it's compatible with black by default
"N801", # class names should use CapWords convention
"N802", # function name should follow snake_case
"N806", # variable in function should follow snake_case
"N812", # lowercase imported as non-lowercase: functional as F
]
[tool.ruff.lint.per-file-ignores]
@@ -168,3 +147,11 @@ markers = [
"benchmark: Benchmark tests",
"slow: Slow tests",
]
[tool.typos.default]
extend-ignore-identifiers-re = [
".*_thw",
".*thw",
"ein",
".*arange"
]
vllm_omni/config/model.py
+1 -1
View File
@@ -2,9 +2,9 @@ from typing import Optional
from pydantic import ConfigDict
from pydantic.dataclasses import dataclass
from vllm.config import ModelConfig, config
import vllm_omni.model_executor.models as me_models
from vllm.config import ModelConfig, config
@config
vllm_omni/core/sched/diffusion_scheduler.py
+14 -45
View File
@@ -4,15 +4,10 @@ from typing import Optional
from vllm.v1.core.kv_cache_manager import KVCacheBlocks
from vllm.v1.core.sched.request_queue import create_request_queue
from vllm.v1.core.sched.scheduler import (
EngineCoreOutputs,
Request,
RequestStatus,
SchedulerOutput,
SpecDecodingStats,
)
from vllm.v1.core.sched.scheduler import EngineCoreOutputs, Request, RequestStatus, SchedulerOutput, SpecDecodingStats
from vllm.v1.core.sched.utils import remove_all
from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput
from vllm_omni.core.sched.output import OmniNewRequestData
from vllm_omni.core.sched.scheduler import OmniScheduler
from vllm_omni.outputs import OmniModelRunnerOutput
@@ -97,17 +92,13 @@ class DiffusionScheduler(OmniScheduler):
num_common_prefix_blocks = [0] * len(self.kv_cache_config.kv_cache_groups)
if self.running:
any_request = self.running[0]
num_common_prefix_blocks = (
self.kv_cache_manager.get_num_common_prefix_blocks(
any_request, len(self.running)
)
num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(
any_request, len(self.running)
)
# Assemble SchedulerOutput
new_reqs_data = [
OmniNewRequestData.from_request(
req, req_to_new_blocks[req.request_id].get_block_ids()
)
OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids())
for req in scheduled_new_reqs
]
cached_reqs_data = self._make_cached_request_data(
@@ -117,9 +108,7 @@ class DiffusionScheduler(OmniScheduler):
scheduled_spec_decode_tokens,
req_to_new_blocks,
)
scheduled_requests = (
scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
)
scheduled_requests = scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs
structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask(
scheduled_requests, scheduled_spec_decode_tokens
)
@@ -177,9 +166,7 @@ class DiffusionScheduler(OmniScheduler):
outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list)
spec_decoding_stats: Optional[SpecDecodingStats] = None
kv_connector_stats = (
kv_connector_output.kv_connector_stats if kv_connector_output else None
)
kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None
# NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more,
# the below loop can be a performance bottleneck. We should do our best
@@ -196,13 +183,9 @@ class DiffusionScheduler(OmniScheduler):
continue
req_index = model_runner_output.req_id_to_index[req_id]
generated_token_ids = (
sampled_token_ids[req_index] if sampled_token_ids else []
)
generated_token_ids = sampled_token_ids[req_index] if sampled_token_ids else []
scheduled_spec_token_ids = (
scheduler_output.scheduled_spec_decode_tokens.get(req_id)
)
scheduled_spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id)
if scheduled_spec_token_ids:
num_draft_tokens = len(scheduled_spec_token_ids)
num_accepted = len(generated_token_ids) - 1
@@ -239,11 +222,7 @@ class DiffusionScheduler(OmniScheduler):
stopped_preempted_reqs.add(request)
# Extract sample logprobs if needed.
if (
request.sampling_params is not None
and request.sampling_params.logprobs is not None
and logprobs
):
if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs:
# NOTE: once we support N tokens per step (spec decode),
# the outer lists can be of length > 1.
new_logprobs = logprobs.slice(req_index, req_index + 1)
@@ -252,9 +231,7 @@ class DiffusionScheduler(OmniScheduler):
# NOTE: structured_output_request should not be None if
# use_structured_output, we have check above, so safe to ignore
# type warning
request.structured_output_request.grammar.accept_tokens( # type: ignore[union-attr] # noqa: E501
req_id, new_token_ids
)
request.structured_output_request.grammar.accept_tokens(req_id, new_token_ids) # type: ignore[union-attr] # noqa: E501
# spec_token_ids comes from the model runner output
if num_nans_in_logits is not None and req_id in num_nans_in_logits:
@@ -263,7 +240,6 @@ class DiffusionScheduler(OmniScheduler):
# Get prompt logprobs for this request.
prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id)
if new_token_ids or pooler_output is not None or kv_transfer_params:
# Add EngineCoreOutput for this Request.
outputs[request.client_index].append(
EngineCoreOutput(
@@ -298,10 +274,7 @@ class DiffusionScheduler(OmniScheduler):
# Create EngineCoreOutputs for all clients that have requests with
# outputs in this step.
engine_core_outputs = {
client_index: EngineCoreOutputs(outputs=outs)
for client_index, outs in outputs.items()
}
engine_core_outputs = {client_index: EngineCoreOutputs(outputs=outs) for client_index, outs in outputs.items()}
finished_req_ids = self.finished_req_ids_dict
if finished_req_ids:
@@ -312,14 +285,10 @@ class DiffusionScheduler(OmniScheduler):
if (eco := engine_core_outputs.get(client_index)) is not None:
eco.finished_requests = finished_set
else:
engine_core_outputs[client_index] = EngineCoreOutputs(
finished_requests=finished_set
)
engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set)
finished_req_ids.clear()
if (
stats := self.make_stats(spec_decoding_stats, kv_connector_stats)
) is not None:
if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None:
# Return stats to only one of the front-ends.
if (eco := next(iter(engine_core_outputs.values()), None)) is None:
# We must return the stats even if there are no request
vllm_omni/core/sched/output.py
+1
View File
@@ -3,6 +3,7 @@ from typing import Optional
from vllm.v1.core.sched.output import NewRequestData
from vllm.v1.request import Request
from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
vllm_omni/core/sched/scheduler.py
+3 -11
View File
@@ -39,22 +39,14 @@ class OmniScheduler(VLLMScheduler):
num_computed_tokens=nr.num_computed_tokens,
lora_request=nr.lora_request,
# Enrich with omni payloads from the live request object
prompt_embeds=(
getattr(request, "prompt_embeds", None) if request else None
),
additional_information=(
getattr(request, "additional_information", None)
if request
else None
),
prompt_embeds=(getattr(request, "prompt_embeds", None) if request else None),
additional_information=(getattr(request, "additional_information", None) if request else None),
)
new_list.append(omni_nr)
scheduler_output.scheduled_new_reqs = new_list # type: ignore[assignment]
except Exception:
# If anything goes wrong, leave the original output unchanged
init_logger(__name__).exception(
"Failed to wrap scheduled_new_reqs with OmniNewRequestData"
)
init_logger(__name__).exception("Failed to wrap scheduled_new_reqs with OmniNewRequestData")
return scheduler_output
vllm_omni/engine/__init__.py
-1
View File
@@ -5,7 +5,6 @@ Engine components for vLLM-omni.
from typing import Any, Optional
import msgspec
from vllm.v1.engine import EngineCoreRequest
vllm_omni/engine/arg_utils.py
+1
View File
@@ -3,6 +3,7 @@ from typing import Optional
from vllm.engine.arg_utils import EngineArgs
from vllm.v1.engine.async_llm import AsyncEngineArgs
from vllm_omni.config import OmniModelConfig
vllm_omni/engine/output_processor.py
+22 -63
View File
@@ -1,7 +1,6 @@
from typing import Any, Callable, Dict, Optional, Union
from typing import Any, Callable, Optional, Union
import torch
from vllm.logger import init_logger
from vllm.outputs import PoolingRequestOutput
from vllm.sampling_params import RequestOutputKind
@@ -10,20 +9,16 @@ from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason
from vllm.v1.engine.detokenizer import IncrementalDetokenizer
from vllm.v1.engine.logprobs import LogprobsProcessor
from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor
from vllm.v1.engine.output_processor import (
OutputProcessorOutput,
RequestOutputCollector,
RequestState,
)
from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestOutputCollector, RequestState
from vllm.v1.engine.parallel_sampling import ParentRequest
from vllm.v1.metrics.stats import IterationStats
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
class OmniRequestState(RequestState):
def __init__(
self,
*args,
@@ -44,7 +39,6 @@ class OmniRequestState(RequestState):
queue: Optional[Any],
log_stats: bool,
) -> "OmniRequestState":
if sampling_params := request.sampling_params:
if not sampling_params.detokenize:
tokenizer = None
@@ -75,9 +69,7 @@ class OmniRequestState(RequestState):
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
),
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,
@@ -93,9 +85,7 @@ class OmniRequestState(RequestState):
log_stats=log_stats,
)
def add_multimodal_tensor(
self, tensor: Optional[torch.Tensor], mm_type: Optional[str]
) -> None:
def add_multimodal_tensor(self, tensor: Optional[torch.Tensor], mm_type: Optional[str]) -> None:
if tensor is None:
return
try:
@@ -137,15 +127,11 @@ class OmniRequestState(RequestState):
if self.parent_req is None:
outputs = [output]
else:
request_id, outputs, finished = self.parent_req.get_outputs(
request_id, output
)
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
)
return self._new_request_output(request_id, outputs, finished, kv_transfer_params)
def _new_completion_output(
self,
@@ -154,9 +140,7 @@ class OmniRequestState(RequestState):
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
)
base_output = super()._new_completion_output(token_ids, finish_reason, stop_reason)
try:
if self.mm_accumulated is not None:
tensor = self.mm_accumulated
@@ -196,14 +180,12 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
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.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
def register_handler(
self, modality: str, handler: Callable[[EngineCoreOutput], None]
) -> None:
def register_handler(self, modality: str, handler: Callable[[EngineCoreOutput], None]) -> None:
self.output_handlers[modality.lower()] = handler
def add_request(
@@ -255,9 +237,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
continue
# 1) Stats
self._update_stats_from_output(
req_state, eco, engine_core_timestamp, iteration_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
@@ -270,9 +250,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
# 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
)
stop_string = req_state.detokenizer.update(new_token_ids, finish_reason == FinishReason.STOP)
if stop_string:
finish_reason = FinishReason.STOP
stop_reason = stop_string
@@ -280,12 +258,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
# 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
):
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(
@@ -310,18 +284,13 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
if ro:
# Attach accumulated multimodal payload if any
try:
if (
isinstance(req_state, OmniRequestState)
and req_state.mm_accumulated is not None
):
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"
)
logger.exception("Error attaching multimodal payload in process_outputs")
if req_state.queue is not None:
req_state.queue.put(ro)
else:
@@ -335,9 +304,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
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
)
self._update_stats_from_finished(req_state, finish_reason, iteration_stats)
if self.tracer:
self.do_tracing(eco, req_state, iteration_stats)
# Cleanup per-request mm state
@@ -352,9 +319,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
# ---- routing helpers ----
def _route_and_normalize(self, eco: EngineCoreOutput) -> None:
output_type = (
getattr(eco, "output_type", self.engine_core_output_type) or ""
).lower()
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:
@@ -385,9 +350,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
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")
)
tensor = self._extract_from_multimodal_outputs(eco, keys=("image", "images", "pixel_values", "pixels"))
if tensor is not None:
eco.pooling_output = tensor
@@ -414,9 +377,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
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")
)
tensor = self._extract_from_multimodal_outputs(eco, keys=("latent", "latents", "z", "posterior"))
if tensor is not None:
eco.pooling_output = tensor
@@ -444,9 +405,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor):
except Exception:
pass
def _extract_from_multimodal_outputs(
self, eco: EngineCoreOutput, keys: tuple[str, ...]
) -> Optional[torch.Tensor]:
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
vllm_omni/engine/processor.py
+11 -38
View File
@@ -3,7 +3,6 @@ from collections.abc import Mapping
from typing import Any, Optional, Union
import torch
from vllm.config import VllmConfig
from vllm.inputs import ProcessorInputs, PromptType
from vllm.inputs.parse import split_enc_dec_inputs
@@ -18,6 +17,7 @@ from vllm.sampling_params import SamplingParams
from vllm.transformers_utils.tokenizer import AnyTokenizer
from vllm.utils import length_from_prompt_token_ids_or_embeds
from vllm.v1.engine.processor import Processor
from vllm_omni.engine import (
AdditionalInformationEntry,
AdditionalInformationPayload,
@@ -78,20 +78,14 @@ class OmniProcessor(Processor):
priority: int = 0,
data_parallel_rank: Optional[int] = None,
) -> tuple[Optional[str], OmniEngineCoreRequest]:
# TODO(woosuk): Support pooling models.
# TODO(woosuk): Support encoder-decoder models.
self._validate_lora(lora_request)
self._validate_params(params)
data_parallel_size = self.vllm_config.parallel_config.data_parallel_size
if data_parallel_rank is not None and not (
0 <= data_parallel_rank < data_parallel_size
):
raise ValueError(
f"data_parallel_rank {data_parallel_rank} "
f"is out of range [0, {data_parallel_size})."
)
if data_parallel_rank is not None and not (0 <= data_parallel_rank < data_parallel_size):
raise ValueError(f"data_parallel_rank {data_parallel_rank} is out of range [0, {data_parallel_size}).")
if arrival_time is None:
arrival_time = time.time()
@@ -142,19 +136,9 @@ class OmniProcessor(Processor):
# discriminated unions of TypedDicts, because of how it handles
# inheritance of TypedDict. If we explicitly extract the items we want
# we can avoid type errors from using `dict.get` later in the method.
prompt_str: Optional[str] = (
None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt")
)
prompt_token_ids = (
decoder_inputs["prompt_token_ids"]
if decoder_inputs["type"] != "embeds"
else None
)
prompt_embeds = (
decoder_inputs["prompt_embeds"]
if decoder_inputs["type"] == "embeds"
else None
)
prompt_str: Optional[str] = None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt")
prompt_token_ids = decoder_inputs["prompt_token_ids"] if decoder_inputs["type"] != "embeds" else None
prompt_embeds = decoder_inputs["prompt_embeds"] if decoder_inputs["type"] == "embeds" else None
sampling_params = None
pooling_params = None
@@ -163,13 +147,9 @@ class OmniProcessor(Processor):
sampling_params = params.clone()
# If unset max tokens, then generate up to the max_model_len.
if sampling_params.max_tokens is None:
seq_len = length_from_prompt_token_ids_or_embeds(
prompt_token_ids, prompt_embeds
)
seq_len = length_from_prompt_token_ids_or_embeds(prompt_token_ids, prompt_embeds)
sampling_params.max_tokens = self.model_config.max_model_len - seq_len
sampling_params.update_from_generation_config(
self.generation_config_fields, eos_token_id
)
sampling_params.update_from_generation_config(self.generation_config_fields, eos_token_id)
if self.tokenizer is not None:
sampling_params.update_from_tokenizer(self.tokenizer)
else:
@@ -204,12 +184,9 @@ class OmniProcessor(Processor):
prompt_embeds_payload: Optional[PromptEmbedsPayload] = None
additional_information_payload: Optional[AdditionalInformationPayload] = None
if "prompt_embeds" in decoder_inputs: # type: ignore[operator]
pe: torch.Tensor = decoder_inputs["prompt_embeds"] # type: ignore[index]
if pe.ndim != 2:
raise ValueError(
"prompt_embeds must be of shape (seq_len, hidden_size)"
)
raise ValueError("prompt_embeds must be of shape (seq_len, hidden_size)")
# Move to CPU and ensure contiguous memory for stable serialization
pe_cpu = pe.detach().to("cpu").contiguous()
seq_len, hidden_size = pe_cpu.shape
@@ -236,13 +213,9 @@ class OmniProcessor(Processor):
elif isinstance(value, list):
entry = AdditionalInformationEntry(list_data=value)
else:
raise ValueError(
"additional_information values must be Tensor or list"
)
raise ValueError("additional_information values must be Tensor or list")
entries[key] = entry
additional_information_payload = AdditionalInformationPayload(
entries=entries
)
additional_information_payload = AdditionalInformationPayload(entries=entries)
return prompt_str, OmniEngineCoreRequest(
request_id=request_id,
vllm_omni/entrypoints/async_omni_llm.py
+21 -57
View File
@@ -4,9 +4,9 @@ import os
import socket
import sys
import time
from collections.abc import AsyncGenerator, Iterable
from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Mapping, Optional, Sequence, Union
from typing import Any, Optional, Union
import torch
@@ -45,11 +45,8 @@ from vllm_omni.entrypoints.log_utils import (
from vllm_omni.entrypoints.omni_stage import OmniStage
from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
from vllm_omni.entrypoints.stage_utils import serialize_obj as _ser
from vllm_omni.entrypoints.utils import (
load_stage_configs_from_model,
load_stage_configs_from_yaml,
)
from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
@@ -68,7 +65,6 @@ class AsyncOmniLLM(EngineClient):
init_timeout: int = 300,
**kwargs,
):
self.batch_timeout = batch_timeout
self._enable_stats: bool = bool(log_stats)
@@ -85,12 +81,8 @@ class AsyncOmniLLM(EngineClient):
remove_old_logs(self._log_file, len(self.stage_list))
configure_orchestrator_logger(logger, self._log_file)
self._stats_file, self._overall_stats_file = init_stats_paths(
self._enable_stats, self._log_file
)
self._initialize_stages(
model, init_sleep_seconds, shm_threshold_bytes, init_timeout
)
self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
def _initialize_stages(
self,
@@ -106,21 +98,14 @@ class AsyncOmniLLM(EngineClient):
idx, cfg = idx_cfg
return idx, OmniStage(cfg)
with ThreadPoolExecutor(
max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))
) as executor:
futures = [
executor.submit(_build_stage, (idx, cfg))
for idx, cfg in enumerate(self.stage_configs)
]
with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor:
futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)]
results: list[tuple[int, OmniStage]] = []
for fut in as_completed(futures):
results.append(fut.result())
results.sort(key=lambda x: x[0])
self.stage_list = [st for _, st in results]
self.default_sampling_params_list = [
st.default_sampling_params for st in self.stage_list
]
self.default_sampling_params_list = [st.default_sampling_params for st in self.stage_list]
logger.debug("[Orchestrator] Loaded %d stages", len(self.stage_list))
self._ctx = mp.get_context("spawn")
@@ -187,9 +172,7 @@ class AsyncOmniLLM(EngineClient):
self,
prompt: PromptType,
request_id: str,
sampling_params_list: Optional[
Union[SamplingParams, Sequence[SamplingParams]]
] = None,
sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
lora_request: Optional[LoRARequest] = None,
trace_headers: Optional[Mapping[str, str]] = None,
priority: int = 0,
@@ -277,13 +260,9 @@ class AsyncOmniLLM(EngineClient):
time.sleep(0.05)
continue
engine_outputs = _load(
result, obj_key="engine_outputs", shm_key="engine_outputs_shm"
)
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
# Mark last output time for this stage whenever we receive outputs
metrics.stage_last_ts[stage_id] = max(
metrics.stage_last_ts[stage_id] or 0.0, time.time()
)
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
try:
_m = result.get("metrics")
if _m is not None:
@@ -314,15 +293,8 @@ class AsyncOmniLLM(EngineClient):
# End-to-end timing and time-per-token for final output
# (only once per request at the designated final stage)
try:
rid_int = (
int(req_id)
if isinstance(req_id, (int, str)) and str(req_id).isdigit()
else req_id
)
if (
stage_id == final_stage_id_for_e2e
and rid_int not in metrics.e2e_done
):
rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
metrics.on_finalize_request(
stage_id,
req_id,
@@ -349,15 +321,13 @@ class AsyncOmniLLM(EngineClient):
next_stage_id = stage_id + 1
if next_stage_id < num_stages:
next_stage: OmniStage = self.stage_list[next_stage_id]
next_inputs = next_stage.process_engine_inputs(
self.stage_list, prompt
)
next_inputs = next_stage.process_engine_inputs(self.stage_list, prompt)
sp_next: SamplingParams = sampling_params_list[next_stage_id]
try:
# Measure transfer size and time (encode + enqueue)
size_bytes = 0
try:
size_bytes = len(_ser(next_inputs))
size_bytes = len(_set(next_inputs))
except Exception:
size_bytes = 0
t0 = time.time()
@@ -654,17 +624,15 @@ class AsyncOmniStageLLM(AsyncLLM):
self.log_stats = log_stats or (stat_loggers is not None)
if not log_stats and stat_loggers is not None:
logger.info(
"AsyncLLM created with log_stats=False and non-empty custom "
"logger list; enabling logging without default stat loggers"
"AsyncLLM created with log_stats=False and non-empty custom logger list; "
"enabling logging without default stat loggers"
)
if self.model_config.skip_tokenizer_init:
self.tokenizer = None
else:
# Tokenizer (+ ensure liveness if running in another process).
self.tokenizer = init_tokenizer_from_configs(
model_config=vllm_config.model_config
)
self.tokenizer = init_tokenizer_from_configs(model_config=vllm_config.model_config)
# Processor (converts Inputs --> EngineCoreRequests).
self.processor = OmniProcessor(
@@ -680,9 +648,7 @@ class AsyncOmniStageLLM(AsyncLLM):
engine_core_output_type=engine_args.engine_output_type,
)
if self.observability_config.otlp_traces_endpoint is not None:
tracer = init_tracer(
"vllm.llm_engine", self.observability_config.otlp_traces_endpoint
)
tracer = init_tracer("vllm.llm_engine", self.observability_config.otlp_traces_endpoint)
self.output_processor.tracer = tracer
# EngineCore (starts the engine in background process).
@@ -736,9 +702,7 @@ class AsyncOmniStageLLM(AsyncLLM):
@classmethod
@deprecate_kwargs(
"disable_log_requests",
additional_message=(
"This argument will have no effect. " "Use `enable_log_requests` instead."
),
additional_message=("This argument will have no effect. Use `enable_log_requests` instead."),
)
def from_vllm_config(
cls,
vllm_omni/entrypoints/cli/main.py
+2 -1
View File
@@ -15,10 +15,11 @@ def main():
vllm_main()
return
else:
import vllm_omni.entrypoints.cli.serve
from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup
from vllm.utils import FlexibleArgumentParser
import vllm_omni.entrypoints.cli.serve
CMD_MODULES = [
vllm_omni.entrypoints.cli.serve,
]
vllm_omni/entrypoints/cli/serve.py
+2 -4
View File
@@ -5,12 +5,12 @@ Omni serve command for vLLM-omni.
import argparse
import uvloop
from vllm.entrypoints.cli.types import CLISubcommand
from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args
from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG
from vllm.logger import init_logger
from vllm.utils import FlexibleArgumentParser
from vllm_omni.entrypoints.openai.api_server import omni_run_server
logger = init_logger(__name__)
@@ -40,9 +40,7 @@ class OmniServeCommand(CLISubcommand):
def validate(self, args: argparse.Namespace) -> None:
validate_parsed_serve_args(args)
def subparser_init(
self, subparsers: argparse._SubParsersAction
) -> FlexibleArgumentParser:
def subparser_init(self, subparsers: argparse._SubParsersAction) -> FlexibleArgumentParser:
serve_parser = subparsers.add_parser(
self.name,
description=DESCRIPTION,
vllm_omni/entrypoints/log_utils.py
+73 -134
View File
@@ -3,12 +3,12 @@ from __future__ import annotations
import logging
import os
import time
from typing import Any, Dict, List, Optional
from typing import Any
from vllm_omni.entrypoints.stage_utils import append_jsonl as _append_jsonl
def remove_old_logs(log_file: Optional[str], num_stages: int) -> None:
def remove_old_logs(log_file: str | None, num_stages: int) -> None:
try:
if not log_file:
return
@@ -48,34 +48,24 @@ def remove_old_logs(log_file: Optional[str], num_stages: int) -> None:
pass
def configure_orchestrator_logger(
logger: logging.Logger, log_file: Optional[str]
) -> None:
def configure_orchestrator_logger(logger: logging.Logger, log_file: str | None) -> None:
try:
if not log_file:
return
has_file_handler = any(
isinstance(h, logging.FileHandler) for h in logger.handlers
)
has_file_handler = any(isinstance(h, logging.FileHandler) for h in logger.handlers)
if not has_file_handler:
fh = logging.FileHandler(log_file)
fh.setLevel(logging.DEBUG)
fh.setFormatter(
logging.Formatter(
"%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s"
)
)
fh.setFormatter(logging.Formatter("%(asctime)s [PID:%(process)d] %(levelname)s: %(message)s"))
logger.addHandler(fh)
logger.setLevel(logging.DEBUG)
except Exception:
pass
def init_stats_paths(
enable_stats: bool, log_file: Optional[str]
) -> tuple[Optional[str], Optional[str]]:
stats_file: Optional[str] = None
overall_file: Optional[str] = None
def init_stats_paths(enable_stats: bool, log_file: str | None) -> tuple[str | None, str | None]:
stats_file: str | None = None
overall_file: str | None = None
try:
if enable_stats and log_file:
stats_file = f"{log_file}.orchestrator.stats.jsonl"
@@ -86,7 +76,7 @@ def init_stats_paths(
return stats_file, overall_file
def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None:
def _safe_append_jsonl(path: str | None, record: dict[str, Any]) -> None:
if not path:
return
try:
@@ -96,7 +86,7 @@ def _safe_append_jsonl(path: Optional[str], record: Dict[str, Any]) -> None:
def log_transfer_tx(
stats_file: Optional[str],
stats_file: str | None,
from_stage: int,
to_stage: int,
request_id: Any,
@@ -120,7 +110,7 @@ def log_transfer_tx(
def log_transfer_rx(
stats_file: Optional[str],
stats_file: str | None,
from_stage: int,
to_stage: int,
request_id: Any,
@@ -139,16 +129,14 @@ def log_transfer_rx(
"rx_decode_time_ms": float(rx_decode_time_ms),
"in_flight_time_ms": float(in_flight_time_ms),
"rx_time_per_kb_ms": (
(float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6))
if rx_bytes > 0
else 0.0
(float(rx_decode_time_ms) / max(float(rx_bytes) / 1024.0, 1e-6)) if rx_bytes > 0 else 0.0
),
},
)
def log_transfer_total(
stats_file: Optional[str],
stats_file: str | None,
from_stage: int,
to_stage: int,
request_id: Any,
@@ -171,16 +159,14 @@ def log_transfer_total(
"rx_decode_time_ms": float(rx_decode_time_ms),
"total_time_ms": float(total_time_ms),
"total_time_per_kb_ms": (
float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6)
if size_bytes > 0
else 0.0
float(total_time_ms) / max(float(size_bytes) / 1024.0, 1e-6) if size_bytes > 0 else 0.0
),
},
)
def log_orchestrator_e2e(
stats_file: Optional[str],
stats_file: str | None,
request_id: Any,
final_stage_id: int,
e2e_time_ms: float,
@@ -200,26 +186,20 @@ def log_orchestrator_e2e(
)
def log_orchestrator_summary(
stats_file: Optional[str], summary: Dict[str, Any]
) -> None:
def log_orchestrator_summary(stats_file: str | None, summary: dict[str, Any]) -> None:
_safe_append_jsonl(stats_file, {"type": "orchestrator_summary", **summary})
def log_overall_summary(
overall_stats_file: Optional[str], summary: Dict[str, Any]
) -> None:
def log_overall_summary(overall_stats_file: str | None, summary: dict[str, Any]) -> None:
_safe_append_jsonl(overall_stats_file, {"type": "overall_summary", **summary})
def log_overall_record(
overall_stats_file: Optional[str], record: Dict[str, Any]
) -> None:
def log_overall_record(overall_stats_file: str | None, record: dict[str, Any]) -> None:
_safe_append_jsonl(overall_stats_file, record)
def log_stage_request_stats(
stats_file: Optional[str],
stats_file: str | None,
stage_id: int,
request_id: Any,
batch_size: int,
@@ -248,7 +228,7 @@ def log_stage_request_stats(
def log_stage_running_avg(
stats_file: Optional[str],
stats_file: str | None,
stage_id: int,
total_tokens: int,
total_gen_time_ms: float,
@@ -267,11 +247,11 @@ def log_stage_running_avg(
def log_stage_batch_stats(
stats_file: Optional[str],
stats_file: str | None,
stage_id: int,
batch_size: int,
batch_gen_time_ms: float,
request_ids: List[Any],
request_ids: list[Any],
) -> None:
_safe_append_jsonl(
stats_file,
@@ -286,23 +266,20 @@ def log_stage_batch_stats(
def compute_and_log_stage_request_stats(
stats_file: Optional[str],
stats_file: str | None,
stage_id: int,
request_id: Any,
batch_size: int,
engine_outputs: List[Any],
engine_outputs: list[Any],
stage_gen_time_ms: float,
rx_transfer_bytes: int,
rx_decode_time_ms: float,
) -> None:
"""Compute per-request metrics and log them in one call."""
num_tokens = count_tokens_from_outputs(engine_outputs)
tokens_per_s = (
(num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
)
tokens_per_s = (num_tokens * 1000.0 / stage_gen_time_ms) if stage_gen_time_ms > 0 else 0.0
rx_mbps = (
(float(rx_transfer_bytes) * 8.0)
/ (max(float(rx_decode_time_ms), 1e-6) * 1000.0)
(float(rx_transfer_bytes) * 8.0) / (max(float(rx_decode_time_ms), 1e-6) * 1000.0)
if rx_transfer_bytes > 0
else 0.0
)
@@ -324,21 +301,19 @@ def compute_and_log_stage_request_stats(
def record_stage_metrics(
per_request: Dict[int, Dict[str, Any]],
stage_req_counts: List[int],
stage_total_time_ms: List[float],
stage_total_tokens: List[int],
per_request: dict[int, dict[str, Any]],
stage_req_counts: list[int],
stage_total_time_ms: list[float],
stage_total_tokens: list[int],
stage_id: int,
req_id: Any,
metrics: Dict[str, Any],
metrics: dict[str, Any],
) -> None:
try:
stage_req_counts[stage_id] += 1
stage_total_tokens[stage_id] += int(metrics.get("num_tokens_out", 0))
rid_int = int(req_id)
pr = per_request.setdefault(
rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}
)
pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
pr_stages = pr["stages"] # type: ignore[index]
pr_stages[stage_id] = {
"stage_gen_time_ms": float(metrics.get("stage_gen_time_ms", 0.0)),
@@ -349,15 +324,15 @@ def record_stage_metrics(
def aggregate_rx_and_maybe_total(
transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
transfer_agg: Dict[tuple[int, int], Dict[str, float]],
per_request: Dict[int, Dict[str, Any]],
transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
transfer_agg: dict[tuple[int, int], dict[str, float]],
per_request: dict[int, dict[str, Any]],
stage_id: int,
req_id: Any,
rx_bytes: float,
rx_ms: float,
in_flight_ms: float,
) -> Optional[tuple[int, float, float]]:
) -> tuple[int, float, float] | None:
try:
# Update RX aggregates for (stage_id-1 -> stage_id)
if stage_id > 0:
@@ -394,9 +369,7 @@ def aggregate_rx_and_maybe_total(
agg["total_count"] += 1.0
# accumulate per-request transfer totals
try:
pr = per_request.setdefault(
rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}
)
pr = per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
pr["transfers_ms"] = float(pr.get("transfers_ms", 0.0)) + total_ms # type: ignore[index]
pr["transfers_bytes"] = int(pr.get("transfers_bytes", 0)) + int(rx_bytes) # type: ignore[index]
except Exception:
@@ -408,8 +381,8 @@ def aggregate_rx_and_maybe_total(
def record_sender_transfer_agg(
transfer_agg: Dict[tuple[int, int], Dict[str, float]],
transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]],
transfer_agg: dict[tuple[int, int], dict[str, float]],
transfer_edge_req: dict[tuple[int, int, int], dict[str, float]],
from_stage: int,
to_stage: int,
req_id: Any,
@@ -447,7 +420,7 @@ def record_sender_transfer_agg(
pass
def count_tokens_from_outputs(engine_outputs: List[Any]) -> int:
def count_tokens_from_outputs(engine_outputs: list[Any]) -> int:
total = 0
for _ro in engine_outputs:
try:
@@ -462,11 +435,11 @@ def count_tokens_from_outputs(engine_outputs: List[Any]) -> int:
def build_stage_summary(
stage_req_counts: List[int],
stage_total_tokens: List[int],
stage_total_time_ms: List[float],
) -> List[Dict[str, Any]]:
summary: List[Dict[str, Any]] = []
stage_req_counts: list[int],
stage_total_tokens: list[int],
stage_total_time_ms: list[float],
) -> list[dict[str, Any]]:
summary: list[dict[str, Any]] = []
for sid in range(len(stage_req_counts)):
reqs = stage_req_counts[sid]
tokens = stage_total_tokens[sid]
@@ -487,31 +460,21 @@ def build_stage_summary(
def build_transfer_summary(
transfer_agg: Dict[tuple[int, int], Dict[str, float]]
) -> List[Dict[str, Any]]:
summary: List[Dict[str, Any]] = []
transfer_agg: dict[tuple[int, int], dict[str, float]],
) -> list[dict[str, Any]]:
summary: list[dict[str, Any]] = []
for (src, dst), agg in transfer_agg.items():
sum_bytes = float(agg.get("sum_bytes", 0.0))
sum_ms = float(agg.get("sum_ms", 0.0))
samples = int(agg.get("count", 0.0))
tx_mbps = (
(sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
)
tx_mbps = (sum_bytes * 8.0) / (max(sum_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
sum_rx_bytes = float(agg.get("sum_rx_bytes", 0.0))
sum_rx_ms = float(agg.get("sum_rx_ms", 0.0))
samples_rx = int(agg.get("rx_count", 0.0))
rx_mbps = (
(sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0)
if sum_rx_bytes > 0
else 0.0
)
rx_mbps = (sum_rx_bytes * 8.0) / (max(sum_rx_ms, 1e-6) * 1000.0) if sum_rx_bytes > 0 else 0.0
sum_total_ms = float(agg.get("sum_total_ms", 0.0))
samples_total = int(agg.get("total_count", 0.0))
total_mbps = (
(sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0)
if sum_bytes > 0
else 0.0
)
total_mbps = (sum_bytes * 8.0) / (max(sum_total_ms, 1e-6) * 1000.0) if sum_bytes > 0 else 0.0
summary.append(
{
"from_stage": src,
@@ -537,40 +500,32 @@ class OrchestratorMetrics:
self,
num_stages: int,
enable_stats: bool,
stats_file: Optional[str],
overall_stats_file: Optional[str],
stats_file: str | None,
overall_stats_file: str | None,
wall_start_ts: float,
) -> None:
self.num_stages = int(num_stages)
self.enable_stats = bool(enable_stats)
self.stats_file = stats_file
self.overall_stats_file = overall_stats_file
self.stage_total_time_ms: List[float] = [0.0 for _ in range(self.num_stages)]
self.stage_total_tokens: List[int] = [0 for _ in range(self.num_stages)]
self.stage_req_counts: List[int] = [0 for _ in range(self.num_stages)]
self.transfer_agg: Dict[tuple[int, int], Dict[str, float]] = {}
self.transfer_edge_req: Dict[tuple[int, int, int], Dict[str, float]] = {}
self.stage_total_time_ms: list[float] = [0.0 for _ in range(self.num_stages)]
self.stage_total_tokens: list[int] = [0 for _ in range(self.num_stages)]
self.stage_req_counts: list[int] = [0 for _ in range(self.num_stages)]
self.transfer_agg: dict[tuple[int, int], dict[str, float]] = {}
self.transfer_edge_req: dict[tuple[int, int, int], dict[str, float]] = {}
self.e2e_total_ms: float = 0.0
self.e2e_total_tokens: int = 0
self.e2e_count: int = 0
self.e2e_done: set[int] = set()
self.per_request: Dict[int, Dict[str, Any]] = {}
self.per_request: dict[int, dict[str, Any]] = {}
self.sum_per_request_transfer_ms: float = 0.0
self.wall_start_ts: float = float(wall_start_ts)
self.last_finish_ts: float = float(wall_start_ts)
self.stage_seen_batches: Dict[int, set] = {
sid: set() for sid in range(self.num_stages)
}
self.stage_first_ts: List[Optional[float]] = [
None for _ in range(self.num_stages)
]
self.stage_last_ts: List[Optional[float]] = [
None for _ in range(self.num_stages)
]
self.stage_seen_batches: dict[int, set] = {sid: set() for sid in range(self.num_stages)}
self.stage_first_ts: list[float | None] = [None for _ in range(self.num_stages)]
self.stage_last_ts: list[float | None] = [None for _ in range(self.num_stages)]
def on_stage_metrics(
self, stage_id: int, req_id: Any, metrics: Dict[str, Any]
) -> None:
def on_stage_metrics(self, stage_id: int, req_id: Any, metrics: dict[str, Any]) -> None:
record_stage_metrics(
self.per_request,
self.stage_req_counts,
@@ -585,9 +540,7 @@ class OrchestratorMetrics:
if batch_id_raw is not None:
batch_id = int(batch_id_raw)
if batch_id not in self.stage_seen_batches[stage_id]:
self.stage_total_time_ms[stage_id] += float(
metrics.get("stage_gen_time_ms", 0.0)
)
self.stage_total_time_ms[stage_id] += float(metrics.get("stage_gen_time_ms", 0.0))
self.stage_seen_batches[stage_id].add(batch_id)
except Exception:
pass
@@ -660,9 +613,7 @@ class OrchestratorMetrics:
float(tx_ms),
)
def on_finalize_request(
self, stage_id: int, req_id: Any, engine_outputs: List[Any], req_start_ts: float
) -> None:
def on_finalize_request(self, stage_id: int, req_id: Any, engine_outputs: list[Any], req_start_ts: float) -> None:
try:
rid_int = int(req_id)
except Exception:
@@ -679,9 +630,7 @@ class OrchestratorMetrics:
self.e2e_total_tokens += int(num_tokens)
self.e2e_count += 1
self.e2e_done.add(rid_int)
pr = self.per_request.setdefault(
rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0}
)
pr = self.per_request.setdefault(rid_int, {"stages": {}, "transfers_ms": 0.0, "transfers_bytes": 0})
per_req_record = {
"type": "overall_request",
"request_id": rid_int,
@@ -696,20 +645,16 @@ class OrchestratorMetrics:
log_overall_record(self.overall_stats_file, per_req_record)
if self.enable_stats and self.stats_file:
e2e_tpt = (e2e_ms / num_tokens) if num_tokens > 0 else 0.0
log_orchestrator_e2e(
self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt
)
log_orchestrator_e2e(self.stats_file, req_id, stage_id, e2e_ms, int(num_tokens), e2e_tpt)
def build_and_log_summary(self, final_stage_id_for_e2e: int) -> Dict[str, Any]:
def build_and_log_summary(self, final_stage_id_for_e2e: int) -> dict[str, Any]:
# Compute stage summary using wall time between first input and last output per stage
stage_summary: List[Dict[str, Any]] = []
stage_summary: list[dict[str, Any]] = []
for sid in range(self.num_stages):
first_ts = self.stage_first_ts[sid]
last_ts = self.stage_last_ts[sid]
total_ms = (
(max(0.0, (last_ts - first_ts)) * 1000.0)
if (first_ts is not None and last_ts is not None)
else 0.0
(max(0.0, (last_ts - first_ts)) * 1000.0) if (first_ts is not None and last_ts is not None) else 0.0
)
reqs = self.stage_req_counts[sid]
tokens = self.stage_total_tokens[sid]
@@ -726,16 +671,10 @@ class OrchestratorMetrics:
}
)
transfer_summary = build_transfer_summary(self.transfer_agg)
e2e_avg_req = (
(self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
)
e2e_avg_tok = (
(self.e2e_total_tokens * 1000.0 / self.e2e_total_ms)
if self.e2e_total_ms > 0
else 0.0
)
e2e_avg_req = (self.e2e_total_ms / self.e2e_count) if self.e2e_count > 0 else 0.0
e2e_avg_tok = (self.e2e_total_tokens * 1000.0 / self.e2e_total_ms) if self.e2e_total_ms > 0 else 0.0
wall_time_ms = max(0.0, (self.last_finish_ts - self.wall_start_ts) * 1000.0)
summary: Dict[str, Any] = {
summary: dict[str, Any] = {
"e2e_requests": int(self.e2e_count),
"e2e_total_time_ms": float(wall_time_ms),
"e2e_sum_time_ms": float(self.e2e_total_ms),
vllm_omni/entrypoints/omni_llm.py
+34 -91
View File
@@ -2,8 +2,9 @@ import multiprocessing as mp
import os
import sys
import time
from collections.abc import Sequence
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Optional, Sequence, Union
from typing import Any, Optional, Union
import cloudpickle
from pydantic import ValidationError
@@ -33,11 +34,8 @@ from vllm_omni.entrypoints.log_utils import (
from vllm_omni.entrypoints.omni_stage import OmniStage
from vllm_omni.entrypoints.stage_utils import encode_for_ipc as _encode
from vllm_omni.entrypoints.stage_utils import maybe_load_from_ipc as _load
from vllm_omni.entrypoints.stage_utils import serialize_obj as _ser
from vllm_omni.entrypoints.utils import (
load_stage_configs_from_model,
load_stage_configs_from_yaml,
)
from vllm_omni.entrypoints.stage_utils import serialize_obj as _set
from vllm_omni.entrypoints.utils import load_stage_configs_from_model, load_stage_configs_from_yaml
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
@@ -77,12 +75,8 @@ class OmniLLM:
remove_old_logs(self._log_file, len(self.stage_list))
configure_orchestrator_logger(logger, self._log_file)
self._stats_file, self._overall_stats_file = init_stats_paths(
self._enable_stats, self._log_file
)
self._initialize_stages(
model, init_sleep_seconds, shm_threshold_bytes, init_timeout
)
self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file)
self._initialize_stages(model, init_sleep_seconds, shm_threshold_bytes, init_timeout)
def _initialize_stages(
self,
@@ -98,13 +92,8 @@ class OmniLLM:
idx, cfg = idx_cfg
return idx, OmniStage(cfg)
with ThreadPoolExecutor(
max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))
) as executor:
futures = [
executor.submit(_build_stage, (idx, cfg))
for idx, cfg in enumerate(self.stage_configs)
]
with ThreadPoolExecutor(max_workers=min(len(self.stage_configs), max(1, os.cpu_count() or 1))) as executor:
futures = [executor.submit(_build_stage, (idx, cfg)) for idx, cfg in enumerate(self.stage_configs)]
results: list[tuple[int, OmniStage]] = []
for fut in as_completed(futures):
results.append(fut.result())
@@ -166,9 +155,7 @@ class OmniLLM:
def generate(
self,
prompts: Union[PromptType, Sequence[PromptType]],
sampling_params_list: Optional[
Union[SamplingParams, Sequence[SamplingParams]]
] = None,
sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
) -> list[OmniRequestOutput]:
try:
return self._run_generation(prompts, sampling_params_list)
@@ -181,19 +168,13 @@ class OmniLLM:
def _run_generation(
self,
prompts: Union[PromptType, Sequence[PromptType]],
sampling_params_list: Optional[
Union[SamplingParams, Sequence[SamplingParams]]
] = None,
sampling_params_list: Optional[Union[SamplingParams, Sequence[SamplingParams]]] = None,
) -> list[OmniRequestOutput]:
logger.debug("[Orchestrator] generate() called")
if sampling_params_list is None:
raise ValueError(
"sampling_params_list is required for pipelined generation"
)
raise ValueError("sampling_params_list is required for pipelined generation")
if len(sampling_params_list) != len(self.stage_list):
raise ValueError(
f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}"
)
raise ValueError(f"Expected {len(self.stage_list)} sampling params, got {len(sampling_params_list)}")
# Normalize prompts to a list for per-request iteration
if not isinstance(prompts, (list, tuple)):
@@ -207,9 +188,7 @@ class OmniLLM:
num_stages = len(self.stage_list)
# Map from request_id to original prompt
request_id_to_prompt: dict[int, PromptType] = {
i: p for i, p in enumerate(request_prompts)
}
request_id_to_prompt: dict[int, PromptType] = {i: p for i, p in enumerate(request_prompts)}
# Track per-request start time for end-to-end timing
_req_start_ts: dict[int, float] = {}
@@ -240,9 +219,7 @@ class OmniLLM:
)
# Seed stage-0 queue with all requests
logger.debug(
"[Orchestrator] Seeding %d requests into stage-0", len(request_prompts)
)
logger.debug("[Orchestrator] Seeding %d requests into stage-0", len(request_prompts))
# Mark first input time for stage-0
metrics.stage_first_ts[0] = metrics.stage_first_ts[0] or time.time()
@@ -287,17 +264,14 @@ class OmniLLM:
continue
if result.get("type") == "stage_ready":
# Only happens when stage is initialized slower than expected, so we wait for a short time and try again
# Only happens when stage is initialized slower than expected,
# so we wait for a short time and try again
time.sleep(0.05)
continue
engine_outputs = _load(
result, obj_key="engine_outputs", shm_key="engine_outputs_shm"
)
engine_outputs = _load(result, obj_key="engine_outputs", shm_key="engine_outputs_shm")
# Mark last output time for this stage whenever we receive outputs
metrics.stage_last_ts[stage_id] = max(
metrics.stage_last_ts[stage_id] or 0.0, time.time()
)
metrics.stage_last_ts[stage_id] = max(metrics.stage_last_ts[stage_id] or 0.0, time.time())
try:
_m = result.get("metrics")
if _m is not None:
@@ -330,17 +304,11 @@ class OmniLLM:
stage_id,
)
# End-to-end timing and time-per-token for final output (only once per request at the designated final stage)
# End-to-end timing and time-per-token for final output
# (only once per request at the designated final stage)
try:
rid_int = (
int(req_id)
if isinstance(req_id, (int, str)) and str(req_id).isdigit()
else req_id
)
if (
stage_id == final_stage_id_for_e2e
and rid_int not in metrics.e2e_done
):
rid_int = int(req_id) if isinstance(req_id, (int, str)) and str(req_id).isdigit() else req_id
if stage_id == final_stage_id_for_e2e and rid_int not in metrics.e2e_done:
metrics.on_finalize_request(
stage_id,
req_id,
@@ -358,15 +326,13 @@ class OmniLLM:
next_stage_id = stage_id + 1
if next_stage_id < num_stages:
next_stage: OmniStage = self.stage_list[next_stage_id]
next_inputs = next_stage.process_engine_inputs(
self.stage_list, [request_id_to_prompt[req_id]]
)
next_inputs = next_stage.process_engine_inputs(self.stage_list, [request_id_to_prompt[req_id]])
sp_next: SamplingParams = sampling_params_list[next_stage_id] # type: ignore[index]
try:
# Measure transfer size and time (encode + enqueue)
size_bytes = 0
try:
size_bytes = len(_ser(next_inputs))
size_bytes = len(_set(next_inputs))
except Exception:
size_bytes = 0
t0 = time.time()
@@ -472,9 +438,7 @@ class OmniLLM:
"Increase initialization wait time (init_sleep_seconds or call-site timeout).",
]
if getattr(self, "_log_file", None):
suggestions.append(
f"Inspect per-stage log files for details: {self._log_file}.stage<id>.log"
)
suggestions.append(f"Inspect per-stage log files for details: {self._log_file}.stage<id>.log")
logger.error(
"[Orchestrator] Stage initialization failed, shutting down. Suggestions:\n- %s",
"\n- ".join(suggestions),
@@ -504,13 +468,9 @@ class OmniStageLLM(LLM):
def __init__(
self,
model: str,
compilation_config: Optional[
Union[int, dict[str, Any], CompilationConfig]
] = None,
compilation_config: Optional[Union[int, dict[str, Any], CompilationConfig]] = None,
hf_overrides: Optional[HfOverrides] = None,
structured_outputs_config: Optional[
Union[dict[str, Any], StructuredOutputsConfig]
] = None,
structured_outputs_config: Optional[Union[dict[str, Any], StructuredOutputsConfig]] = None,
**kwargs,
):
"""LLM constructor."""
@@ -524,9 +484,7 @@ class OmniStageLLM(LLM):
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
):
if "kv_transfer_config" in kwargs and isinstance(kwargs["kv_transfer_config"], dict):
from vllm.config.kv_transfer import KVTransferConfig
raw_config_dict = kwargs["kv_transfer_config"]
@@ -534,8 +492,7 @@ class OmniStageLLM(LLM):
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",
"Failed to convert 'kv_transfer_config' dict to KVTransferConfig object. Dict: %s. Error: %s",
raw_config_dict,
e,
)
@@ -548,16 +505,10 @@ class OmniStageLLM(LLM):
if compilation_config is not None:
if isinstance(compilation_config, int):
compilation_config_instance = CompilationConfig(
level=compilation_config
)
compilation_config_instance = CompilationConfig(level=compilation_config)
elif isinstance(compilation_config, dict):
compilation_config_instance = CompilationConfig(
**{
k: v
for k, v in compilation_config.items()
if is_init_field(CompilationConfig, k)
}
**{k: v for k, v in compilation_config.items() if is_init_field(CompilationConfig, k)}
)
else:
compilation_config_instance = compilation_config
@@ -567,11 +518,7 @@ class OmniStageLLM(LLM):
if structured_outputs_config is not None:
if isinstance(structured_outputs_config, dict):
structured_outputs_instance = StructuredOutputsConfig(
**{
k: v
for k, v in structured_outputs_config.items()
if is_init_field(StructuredOutputsConfig, k)
}
**{k: v for k, v in structured_outputs_config.items() if is_init_field(StructuredOutputsConfig, k)}
)
else:
structured_outputs_instance = structured_outputs_config
@@ -587,9 +534,7 @@ class OmniStageLLM(LLM):
)
# 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 = 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,
@@ -611,6 +556,4 @@ class OmniStageLLM(LLM):
# Load the Input/Output processor plugin if any
io_processor_plugin = self.llm_engine.model_config.io_processor_plugin
self.io_processor = get_io_processor(
self.llm_engine.vllm_config, io_processor_plugin
)
self.io_processor = get_io_processor(self.llm_engine.vllm_config, io_processor_plugin)
vllm_omni/entrypoints/omni_stage.py
+56 -129
View File
@@ -14,7 +14,7 @@ import asyncio
import importlib
import logging
import multiprocessing as mp
from typing import Any, Dict, List, Optional, Union
from typing import Any, Optional, Union
from vllm.inputs import TextPrompt
from vllm.inputs.preprocess import InputPreprocessor
@@ -25,6 +25,7 @@ from vllm.usage.usage_lib import UsageContext
from vllm.v1.engine import EngineCoreOutput
from vllm.v1.engine.async_llm import AsyncLLM
from vllm.v1.engine.llm_engine import LLMEngine
from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs
from vllm_omni.entrypoints.stage_utils import (
_to_dict,
@@ -55,9 +56,7 @@ class OmniStage:
self.is_comprehension = getattr(stage_config, "is_comprehension", False)
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_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:
@@ -66,9 +65,7 @@ class OmniStage:
self.final_output = getattr(stage_config, "final_output", False)
self.final_output_type = getattr(stage_config, "final_output_type", None)
default_sampling_params = getattr(stage_config, "default_sampling_params", {})
self.default_sampling_params = SamplingParams(
**_to_dict(default_sampling_params)
)
self.default_sampling_params = SamplingParams(**_to_dict(default_sampling_params))
# Runtime orchestration state (added)
self._in_q: Optional[mp.Queue] = None
self._out_q: Optional[mp.Queue] = None
@@ -120,16 +117,14 @@ class OmniStage:
ctx: Optional[mp.context.BaseContext] = None,
batch_timeout: int = 10,
) -> None:
assert (
self._in_q is not None and self._out_q is not None
), "Queues must be attached before start_process"
assert self._in_q is not None and self._out_q is not None, "Queues must be attached before start_process"
self._log_file = log_file
self._shm_threshold_bytes = shm_threshold_bytes
ctx = ctx or mp.get_context("spawn")
# Prepare lightweight dict config for worker
engine_args = _to_dict(self.engine_args)
runtime_cfg = _to_dict(getattr(self.stage_config, "runtime", {}))
stage_payload: Dict[str, Any] = {
stage_payload: dict[str, Any] = {
"stage_id": self.stage_id,
"engine_args": engine_args,
"runtime": runtime_cfg,
@@ -167,29 +162,23 @@ class OmniStage:
try:
self._in_q.put_nowait(None)
except Exception as e:
self._logger.warning(
"[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e
)
self._logger.warning("[Stage-%s] Failed to send shutdown to in_q: %s", self.stage_id, e)
if self._proc is not None:
try:
self._proc.join(timeout=5)
except Exception as e:
self._logger.debug(
"[Stage-%s] join() failed: %s", self.stage_id, e, exc_info=True
)
self._logger.debug("[Stage-%s] join() failed: %s", self.stage_id, e, exc_info=True)
if self._proc.is_alive():
try:
self._proc.terminate()
except Exception as e:
self._logger.warning(
"[Stage-%s] terminate() failed: %s", self.stage_id, e
)
self._logger.warning("[Stage-%s] terminate() failed: %s", self.stage_id, e)
def submit(self, payload: Dict[str, Any]) -> None:
def submit(self, payload: dict[str, Any]) -> None:
assert self._in_q is not None
self._in_q.put(payload)
def try_collect(self) -> Optional[Dict[str, Any]]:
def try_collect(self) -> Optional[dict[str, Any]]:
assert self._out_q is not None
try:
return self._out_q.get_nowait()
@@ -198,7 +187,7 @@ class OmniStage:
def process_engine_inputs(
self, stage_list, prompt: Union[OmniTokensPrompt, TextPrompt] = None
) -> List[Union[OmniTokensPrompt, TextPrompt]]:
) -> list[Union[OmniTokensPrompt, TextPrompt]]:
"""Process the engine input for the stage."""
if self.custom_process_input_func is None:
engine_inputs = []
@@ -216,25 +205,19 @@ class OmniStage:
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
),
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
)
return self.custom_process_input_func(stage_list, engine_input_source, prompt)
def _stage_worker(
model: str,
stage_payload: Dict[str, Any],
stage_payload: dict[str, Any],
in_q: mp.Queue,
out_q: mp.Queue,
log_file: Optional[str] = None,
@@ -267,9 +250,7 @@ def _stage_worker(
fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
fh.setLevel(_logging.DEBUG)
fh.setFormatter(
_logging.Formatter(
"%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s"
)
_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")
)
class _StageFilter(_logging.Filter):
@@ -295,9 +276,7 @@ def _stage_worker(
try:
set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
except Exception as e:
_logging.getLogger(__name__).warning(
"[Stage-%s] Device setup failed: %s", stage_id, e
)
_logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
# Init LLM
_logging.getLogger(__name__).debug(
@@ -318,13 +297,11 @@ def _stage_worker(
task = in_q.get()
_recv_dequeue_ts = _time.time()
if task is None:
_logging.getLogger(__name__).debug(
"[Stage-%s] Received shutdown signal", stage_id
)
_logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id)
break
max_batch_size = int(runtime_cfg.get("max_batch_size", 1) or 1)
batch_tasks: List[Dict[str, Any]] = [task]
batch_tasks: list[dict[str, Any]] = [task]
if max_batch_size > 1:
while len(batch_tasks) < max_batch_size:
if not in_q.empty():
@@ -336,11 +313,11 @@ def _stage_worker(
else:
break
batch_request_ids: List[Any] = []
batch_engine_inputs: List[Any] = []
_rx_bytes_by_rid: Dict[Any, int] = {}
_rx_decode_ms_by_rid: Dict[Any, float] = {}
_in_flight_ms_by_rid: Dict[Any, float] = {}
batch_request_ids: list[Any] = []
batch_engine_inputs: list[Any] = []
_rx_bytes_by_rid: dict[Any, int] = {}
_rx_decode_ms_by_rid: dict[Any, float] = {}
_in_flight_ms_by_rid: dict[Any, float] = {}
for t in batch_tasks:
rid = t["request_id"]
try:
@@ -351,9 +328,7 @@ def _stage_worker(
_in_flight_ms_by_rid[rid] = 0.0
except Exception:
_in_flight_ms_by_rid[rid] = 0.0
ein, _rx_metrics = maybe_load_from_ipc_with_metrics(
t, obj_key="engine_inputs", shm_key="engine_inputs_shm"
)
ein, _rx_metrics = maybe_load_from_ipc_with_metrics(t, obj_key="engine_inputs", shm_key="engine_inputs_shm")
_rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
_rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
batch_request_ids.append(rid)
@@ -362,9 +337,7 @@ def _stage_worker(
elif isinstance(ein, dict):
batch_engine_inputs.append(ein)
else:
_logging.getLogger(__name__).exception(
"[Stage-%s] Invalid engine input type: %s", stage_id, type(ein)
)
_logging.getLogger(__name__).exception("[Stage-%s] Invalid engine input type: %s", stage_id, type(ein))
sampling_params = batch_tasks[0]["sampling_params"]
_logging.getLogger(__name__).debug(
"[Stage-%s] Received batch size=%d, request_ids=%s",
@@ -380,20 +353,16 @@ def _stage_worker(
print("--------------------------------", flush=True)
try:
_batch_seq += 1
gen_outputs: List[Any] = []
gen_outputs: list[Any] = []
_gen_t0 = _time.time()
for ro in stage_engine.generate(
batch_engine_inputs, sampling_params, use_tqdm=False
):
for ro in stage_engine.generate(batch_engine_inputs, sampling_params, use_tqdm=False):
gen_outputs.append(ro)
_gen_t1 = _time.time()
_gen_ms = (_gen_t1 - _gen_t0) * 1000.0
# Group outputs per request id with fallback
req_to_outputs: Dict[Any, List[Any]] = {
rid: [] for rid in batch_request_ids
}
unmapped: List[Any] = []
req_to_outputs: dict[Any, list[Any]] = {rid: [] for rid in batch_request_ids}
unmapped: list[Any] = []
for ro in gen_outputs:
rid = getattr(ro, "request_id", None)
if rid in req_to_outputs:
@@ -416,9 +385,7 @@ def _stage_worker(
if _stats_file:
_avg_tokens_per_s = (
(_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms)
if _agg_total_gen_time_ms > 0
else 0.0
(_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
)
log_stage_running_avg(
_stats_file,
@@ -446,9 +413,7 @@ def _stage_worker(
"batch_id": int(_batch_seq),
"rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
"rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
"rx_in_flight_time_ms": float(
_in_flight_ms_by_rid.get(rid, 0.0)
),
"rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
}
if _stats_file:
compute_and_log_stage_request_stats(
@@ -486,17 +451,11 @@ def _stage_worker(
"stage_id": stage_id,
"engine_outputs": r_outputs,
"metrics": {
"num_tokens_out": int(
count_tokens_from_outputs(r_outputs)
),
"num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
"stage_gen_time_ms": _gen_ms,
"rx_decode_time_ms": float(
_rx_decode_ms_by_rid.get(rid, 0.0)
),
"rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
"rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
"rx_in_flight_time_ms": float(
_in_flight_ms_by_rid.get(rid, 0.0)
),
"rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
},
}
)
@@ -506,9 +465,7 @@ def _stage_worker(
rid,
)
except Exception as e:
_logging.getLogger(__name__).exception(
"[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e
)
_logging.getLogger(__name__).exception("[Stage-%s] Failed on batch %s: %s", stage_id, batch_request_ids, e)
for rid in batch_request_ids:
out_q.put(
{
@@ -522,23 +479,19 @@ def _stage_worker(
def _stage_worker_async_entry(
omni_stage: OmniStage,
model: str,
stage_payload: Dict[str, Any],
stage_payload: dict[str, Any],
in_q: mp.Queue,
out_q: mp.Queue,
log_file: Optional[str] = None,
batch_timeout: int = 10,
) -> None:
asyncio.run(
_stage_worker_async(
omni_stage, model, stage_payload, in_q, out_q, log_file, batch_timeout
)
)
asyncio.run(_stage_worker_async(omni_stage, model, stage_payload, in_q, out_q, log_file, batch_timeout))
async def _stage_worker_async(
omni_stage: OmniStage,
model: str,
stage_payload: Dict[str, Any],
stage_payload: dict[str, Any],
in_q: mp.Queue,
out_q: mp.Queue,
log_file: Optional[str] = None,
@@ -571,10 +524,8 @@ async def _stage_worker_async(
fh = _logging.FileHandler(f"{log_file}.stage{stage_id}.log")
fh.setLevel(_logging.DEBUG)
fh.setFormatter(
_logging.Formatter(
"%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s" # noqa: E501
)
)
_logging.Formatter("%(asctime)s [PID:%(process)d] [Stage-%(stage)s] %(levelname)s: %(message)s")
) # noqa: E501
class _StageFilter(_logging.Filter):
def filter(self, record: _logging.LogRecord) -> bool:
@@ -600,9 +551,7 @@ async def _stage_worker_async(
try:
set_stage_gpu_devices(stage_id, runtime_cfg.get("devices"))
except Exception as e:
_logging.getLogger(__name__).warning(
"[Stage-%s] Device setup failed: %s", stage_id, e
)
_logging.getLogger(__name__).warning("[Stage-%s] Device setup failed: %s", stage_id, e)
# Init LLM
_logging.getLogger(__name__).debug(
@@ -640,49 +589,37 @@ async def _stage_worker_async(
}
)
except Exception as e:
_logging.getLogger(__name__).warning(
"[Stage-%s] Failed to send stage ready signal: %s", stage_id, e
)
_logging.getLogger(__name__).warning("[Stage-%s] Failed to send stage ready signal: %s", stage_id, e)
# Batch processing loop
while True:
task = in_q.get()
_recv_dequeue_ts = _time.time()
if task is None:
_logging.getLogger(__name__).debug(
"[Stage-%s] Received shutdown signal", stage_id
)
_logging.getLogger(__name__).debug("[Stage-%s] Received shutdown signal", stage_id)
break
_rx_bytes_by_rid: Dict[Any, int] = {}
_rx_decode_ms_by_rid: Dict[Any, float] = {}
_in_flight_ms_by_rid: Dict[Any, float] = {}
_rx_bytes_by_rid: dict[Any, int] = {}
_rx_decode_ms_by_rid: dict[Any, float] = {}
_in_flight_ms_by_rid: dict[Any, float] = {}
rid = task["request_id"]
try:
sent_ts = (
float(task.get("sent_ts", None)) if isinstance(task, dict) else None
)
sent_ts = float(task.get("sent_ts", None)) if isinstance(task, dict) else None
if sent_ts is not None:
_in_flight_ms_by_rid[rid] = (_recv_dequeue_ts - sent_ts) * 1000.0
else:
_in_flight_ms_by_rid[rid] = 0.0
except Exception:
_in_flight_ms_by_rid[rid] = 0.0
ein, _rx_metrics = maybe_load_from_ipc_with_metrics(
task, obj_key="engine_inputs", shm_key="engine_inputs_shm"
)
ein, _rx_metrics = maybe_load_from_ipc_with_metrics(task, obj_key="engine_inputs", shm_key="engine_inputs_shm")
_rx_decode_ms_by_rid[rid] = float(_rx_metrics.get("rx_decode_time_ms", 0.0))
_rx_bytes_by_rid[rid] = int(_rx_metrics.get("rx_transfer_bytes", 0))
sampling_params = task["sampling_params"]
_logging.getLogger(__name__).debug(
"[Stage-%s] Received batch size=1, request_ids=%d", stage_id, rid
)
_logging.getLogger(__name__).debug("[Stage-%s] Received batch size=1, request_ids=%d", stage_id, rid)
print("--------------------------------", flush=True)
print(
f"[Stage-{stage_id}] Received batch size=1, request_ids={rid}", flush=True
)
print(f"[Stage-{stage_id}] Received batch size=1, request_ids={rid}", flush=True)
print("--------------------------------", flush=True)
try:
_batch_seq += 1
@@ -702,9 +639,7 @@ async def _stage_worker_async(
if _stats_file:
_avg_tokens_per_s = (
(_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms)
if _agg_total_gen_time_ms > 0
else 0.0
(_agg_total_tokens * 1000.0 / _agg_total_gen_time_ms) if _agg_total_gen_time_ms > 0 else 0.0
)
log_stage_running_avg(
_stats_file,
@@ -769,24 +704,16 @@ async def _stage_worker_async(
"metrics": {
"num_tokens_out": int(count_tokens_from_outputs(r_outputs)),
"stage_gen_time_ms": _gen_ms,
"rx_decode_time_ms": float(
_rx_decode_ms_by_rid.get(rid, 0.0)
),
"rx_decode_time_ms": float(_rx_decode_ms_by_rid.get(rid, 0.0)),
"rx_transfer_bytes": int(_rx_bytes_by_rid.get(rid, 0)),
"rx_in_flight_time_ms": float(
_in_flight_ms_by_rid.get(rid, 0.0)
),
"rx_in_flight_time_ms": float(_in_flight_ms_by_rid.get(rid, 0.0)),
},
}
)
_logging.getLogger(__name__).debug(
"[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid
)
_logging.getLogger(__name__).debug("[Stage-%s] Enqueued result for request %s to downstream", stage_id, rid)
except Exception as e:
_logging.getLogger(__name__).exception(
"[Stage-%s] Failed on request %s: %s", stage_id, rid, e
)
_logging.getLogger(__name__).exception("[Stage-%s] Failed on request %s: %s", stage_id, rid, e)
out_q.put(
{
"request_id": rid,
vllm_omni/entrypoints/openai/api_server.py
+16 -49
View File
@@ -2,22 +2,18 @@ import multiprocessing
import multiprocessing.forkserver as forkserver
import os
from argparse import Namespace
from collections.abc import AsyncIterator
from contextlib import asynccontextmanager
from http import HTTPStatus
from typing import Any, AsyncIterator, Optional
from typing import Any, Optional
import vllm.envs as envs
from fastapi import Depends, HTTPException, Request
from fastapi.responses import JSONResponse, StreamingResponse
from starlette.datastructures import State
import vllm.envs as envs
from vllm.config import VllmConfig
from vllm.engine.protocol import EngineClient
from vllm.entrypoints.chat_utils import (
load_chat_template,
resolve_hf_chat_template,
resolve_mistral_chat_template,
)
from vllm.entrypoints.chat_utils import load_chat_template, resolve_hf_chat_template, resolve_mistral_chat_template
from vllm.entrypoints.launcher import serve_http
from vllm.entrypoints.logger import RequestLogger
from vllm.entrypoints.openai.api_server import (
@@ -29,16 +25,8 @@ from vllm.entrypoints.openai.api_server import (
setup_server,
validate_json_request,
)
from vllm.entrypoints.openai.protocol import (
ChatCompletionRequest,
ChatCompletionResponse,
ErrorResponse,
)
from vllm.entrypoints.openai.serving_models import (
BaseModelPath,
LoRAModulePath,
OpenAIServingModels,
)
from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ChatCompletionResponse, ErrorResponse
from vllm.entrypoints.openai.serving_models import BaseModelPath, LoRAModulePath, OpenAIServingModels
from vllm.entrypoints.openai.tool_parsers import ToolParserManager
# yapf conflicts with isort for this block
@@ -49,6 +37,7 @@ from vllm.entrypoints.utils import load_aware_call, with_cancellation
from vllm.logger import init_logger
from vllm.transformers_utils.tokenizer import MistralTokenizer
from vllm.utils import decorate_logs
from vllm_omni.entrypoints.async_omni_llm import AsyncOmniLLM
from vllm_omni.entrypoints.openai.serving_chat import OmniOpenAIServingChat
@@ -65,9 +54,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None:
await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs)
async def omni_run_server_worker(
listen_address, sock, args, client_config=None, **uvicorn_kwargs
) -> None:
async def omni_run_server_worker(listen_address, sock, args, client_config=None, **uvicorn_kwargs) -> None:
"""Run a single API server worker."""
if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3:
@@ -127,7 +114,6 @@ async def build_async_omni_llm(
disable_frontend_multiprocessing: Optional[bool] = None,
client_config: Optional[dict[str, Any]] = None,
) -> AsyncIterator[EngineClient]:
if os.getenv("VLLM_WORKER_MULTIPROC_METHOD") == "forkserver":
# The executor is expected to be mp.
# Pre-import heavy modules in the forkserver process
@@ -174,9 +160,7 @@ async def build_async_omni_llm_from_stage_config(
try:
if getattr(args, "stage_configs_path", None):
async_omni_llm = AsyncOmniLLM(
model=args.model, stage_configs_path=args.stage_configs_path
)
async_omni_llm = AsyncOmniLLM(model=args.model, stage_configs_path=args.stage_configs_path)
else:
async_omni_llm = AsyncOmniLLM(model=args.model)
@@ -195,7 +179,6 @@ async def omni_init_app_state(
state: State,
args: Namespace,
) -> None:
if args.served_model_name is not None:
served_model_names = args.served_model_name
else:
@@ -206,9 +189,7 @@ async def omni_init_app_state(
else:
request_logger = None
base_model_paths = [
BaseModelPath(name=name, model_path=args.model) for name in served_model_names
]
base_model_paths = [BaseModelPath(name=name, model_path=args.model) for name in served_model_names]
state.engine_client = engine_client
state.log_stats = not args.disable_log_stats
state.vllm_config = vllm_config
@@ -225,9 +206,7 @@ async def omni_init_app_state(
if isinstance(tokenizer, MistralTokenizer):
# The warning is logged in resolve_mistral_chat_template.
resolved_chat_template = resolve_mistral_chat_template(
chat_template=resolved_chat_template
)
resolved_chat_template = resolve_mistral_chat_template(chat_template=resolved_chat_template)
else:
hf_chat_template = resolve_hf_chat_template(
tokenizer=tokenizer,
@@ -238,9 +217,7 @@ async def omni_init_app_state(
if hf_chat_template != resolved_chat_template:
logger.warning(
"Using supplied chat template: %s\n"
"It is different from official chat template '%s'. "
"This discrepancy may lead to performance degradation.",
"Using supplied chat template: %s\nIt is different from official chat template '%s'. This discrepancy may lead to performance degradation.", # noqa: E501
resolved_chat_template,
args.model,
)
@@ -256,11 +233,7 @@ async def omni_init_app_state(
tool_server = None
# Merge default_mm_loras into the static lora_modules
default_mm_loras = (
vllm_config.lora_config.default_mm_loras
if vllm_config.lora_config is not None
else {}
)
default_mm_loras = vllm_config.lora_config.default_mm_loras if vllm_config.lora_config is not None else {}
lora_modules = args.lora_modules
if default_mm_loras:
@@ -326,19 +299,13 @@ def Omnichat(request: Request) -> Optional[OmniOpenAIServingChat]:
async def create_chat_completion(request: ChatCompletionRequest, raw_request: Request):
handler = Omnichat(raw_request)
if handler is None:
return base(raw_request).create_error_response(
message="The model does not support Chat Completions API"
)
return base(raw_request).create_error_response(message="The model does not support Chat Completions API")
try:
generator = await handler.create_chat_completion(request, raw_request)
except Exception as e:
raise HTTPException(
status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
) from e
raise HTTPException(status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)) from e
if isinstance(generator, ErrorResponse):
return JSONResponse(
content=generator.model_dump(), status_code=generator.error.code
)
return JSONResponse(content=generator.model_dump(), status_code=generator.error.code)
elif isinstance(generator, ChatCompletionResponse):
return JSONResponse(content=generator.model_dump())
vllm_omni/entrypoints/openai/serving_chat.py
+36 -120
View File
@@ -17,15 +17,8 @@ try:
except ImportError:
soundfile = None
from openai.types.chat.chat_completion_audio import (
ChatCompletionAudio as OpenAIChatCompletionAudio,
)
from vllm.entrypoints.chat_utils import (
ConversationMessage,
get_history_tool_calls_cnt,
make_tool_call_id,
)
from openai.types.chat.chat_completion_audio import ChatCompletionAudio as OpenAIChatCompletionAudio
from vllm.entrypoints.chat_utils import ConversationMessage, get_history_tool_calls_cnt, make_tool_call_id
from vllm.entrypoints.harmony_utils import parse_chat_output
from vllm.entrypoints.openai.protocol import (
ChatCompletionNamedToolChoiceParam,
@@ -56,6 +49,7 @@ from vllm.transformers_utils.tokenizers import (
validate_request_params,
)
from vllm.utils import as_list
from vllm_omni.outputs import OmniRequestOutput
logger = init_logger(__name__)
@@ -86,9 +80,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
raise self.engine_client.dead_error
try:
lora_request = self._maybe_get_adapters(
request, supports_default_mm_loras=True
)
lora_request = self._maybe_get_adapters(request, supports_default_mm_loras=True)
model_name = self.models.model_name(lora_request)
@@ -113,14 +105,10 @@ class OmniOpenAIServingChat(OpenAIServingChat):
# for hf tokenizers, "auto" tools requires
# --enable-auto-tool-choice and --tool-call-parser
return self.create_error_response(
'"auto" tool choice requires '
"--enable-auto-tool-choice and --tool-call-parser to be set"
'"auto" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set'
)
if request.tools is None or (
request.tool_choice == "none"
and self.exclude_tools_when_tool_choice_none
):
if request.tools is None or (request.tool_choice == "none" and self.exclude_tools_when_tool_choice_none):
tool_dicts = None
else:
tool_dicts = [tool.model_dump() for tool in request.tools]
@@ -130,14 +118,10 @@ class OmniOpenAIServingChat(OpenAIServingChat):
chat_template_kwargs = request.chat_template_kwargs
if not self.trust_request_chat_template and (
request_chat_template is not None
or (
chat_template_kwargs
and chat_template_kwargs.get("chat_template") is not None
)
or (chat_template_kwargs and chat_template_kwargs.get("chat_template") is not None)
):
return self.create_error_response(
"Chat template is passed with request, but "
"--trust-request-chat-template is not set. "
"Chat template is passed with request, but --trust-request-chat-template is not set. "
"Refused request with untrusted chat template."
)
(
@@ -163,9 +147,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
logger.exception("Error in preprocessing prompt inputs")
return self.create_error_response(f"{e} {e.__cause__}")
request_id = (
"chatcmpl-" f"{self._base_request_id(raw_request, request.request_id)}"
)
request_id = f"chatcmpl-{self._base_request_id(raw_request, request.request_id)}"
request_metadata = RequestResponseMetadata(request_id=request_id)
if raw_request:
@@ -175,11 +157,8 @@ class OmniOpenAIServingChat(OpenAIServingChat):
generators: list[AsyncGenerator[RequestOutput, None]] = []
try:
for i, engine_prompt in enumerate(engine_prompts):
if hasattr(request, "sampling_params_list"):
sampling_params_list = self._to_sampling_params_list(
request.sampling_params_list
)
sampling_params_list = self._to_sampling_params_list(request.sampling_params_list)
else:
sampling_params_list = None
@@ -190,11 +169,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
lora_request=lora_request,
)
trace_headers = (
None
if raw_request is None
else await self._get_trace_headers(raw_request.headers)
)
trace_headers = None if raw_request is None else await self._get_trace_headers(raw_request.headers)
generator = self.engine_client.generate(
prompt=engine_prompt,
@@ -240,10 +215,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
# TODO: Use a vllm-specific Validation Error
return self.create_error_response(str(e))
def _to_sampling_params_list(
self, sampling_params_list: list[dict]
) -> list[SamplingParams]:
def _to_sampling_params_list(self, sampling_params_list: list[dict]) -> list[SamplingParams]:
final_sampling_params_list = []
for sampling_params in sampling_params_list:
if isinstance(sampling_params, dict):
@@ -273,10 +245,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
prompt_token_ids = getattr(inputs, "prompt_token_ids", None)
logger.info(
"Received request %s: prompt: %r, "
"params_list: %s, prompt_token_ids: %s, "
"prompt_embeds shape: %s, "
"lora_request: %s.",
"Received request %s: prompt: %r, params_list: %s, prompt_token_ids: %s, prompt_embeds shape: %s, lora_request: %s.", # noqa: E501
request_id,
prompt,
params_list,
@@ -295,7 +264,6 @@ class OmniOpenAIServingChat(OpenAIServingChat):
tokenizer: AnyTokenizer,
request_metadata: RequestResponseMetadata,
) -> Union[ErrorResponse, ChatCompletionResponse]:
created_time = int(time.time())
final_res: Optional[RequestOutput] = None
@@ -328,15 +296,11 @@ class OmniOpenAIServingChat(OpenAIServingChat):
prompt_logprobs,
prompt_token_ids,
kv_transfer_params,
) = self._create_text_choice(
request, omni_outputs, tokenizer, conversation, role
)
) = self._create_text_choice(request, omni_outputs, tokenizer, conversation, role)
elif omni_outputs.final_output_type == "audio":
choices_data = self._create_audio_choice(omni_outputs, role)
else:
logger.warning(
f"Unsupported final output type: {omni_outputs.final_output_type}"
)
logger.warning(f"Unsupported final output type: {omni_outputs.final_output_type}")
continue
choices.extend(choices_data)
@@ -361,12 +325,8 @@ class OmniOpenAIServingChat(OpenAIServingChat):
# For tool calls, log the function name and arguments
tool_call_descriptions = []
for tc in choice.message.tool_calls:
if hasattr(tc.function, "name") and hasattr(
tc.function, "arguments"
):
tool_call_descriptions.append(
f"{tc.function.name}({tc.function.arguments})"
)
if hasattr(tc.function, "name") and hasattr(tc.function, "arguments"):
tool_call_descriptions.append(f"{tc.function.name}({tc.function.arguments})")
tool_calls_str = ", ".join(tool_call_descriptions)
output_text = f"[tool_calls: {tool_calls_str}]"
@@ -469,9 +429,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
return self.create_error_response(str(e))
# If the reasoning parser is enabled,
# tool calls are extracted exclusively from the content.
reasoning_content, content = reasoning_parser.extract_reasoning_content(
output.text, request=request
)
reasoning_content, content = reasoning_parser.extract_reasoning_content(output.text, request=request)
if not request.include_reasoning:
reasoning_content = None
else:
@@ -485,21 +443,11 @@ class OmniOpenAIServingChat(OpenAIServingChat):
not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam)
and request.tool_choice != "required"
):
message = ChatMessage(
role=role, reasoning_content=reasoning_content, content=content
)
message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
# if the request uses tools and specified a tool choice
elif (
request.tool_choice
and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam
):
tool_call_class = (
MistralToolCall
if isinstance(tokenizer, MistralTokenizer)
else ToolCall
)
elif request.tool_choice and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam:
tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall
message = ChatMessage(
role=role,
reasoning_content=reasoning_content,
@@ -515,18 +463,12 @@ class OmniOpenAIServingChat(OpenAIServingChat):
)
elif request.tool_choice and request.tool_choice == "required":
tool_call_class = (
MistralToolCall
if isinstance(tokenizer, MistralTokenizer)
else ToolCall
)
tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall
# the fields of FunctionDefinition are a superset of the
# tool call outputs and can be used for parsing
assert content is not None
tool_calls = TypeAdapter(list[FunctionDefinition]).validate_json(
content
)
tool_calls = TypeAdapter(list[FunctionDefinition]).validate_json(content)
tool_call_ids = []
for tool_call in tool_calls:
tool_call_ids.append(
@@ -545,9 +487,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
id=tool_call_ids[i],
function=FunctionCall(
name=tool_call.name,
arguments=json.dumps(
tool_call.parameters, ensure_ascii=False
),
arguments=json.dumps(tool_call.parameters, ensure_ascii=False),
),
)
for i, tool_call in enumerate(tool_calls)
@@ -558,10 +498,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
# if the request doesn't use tool choice
# OR specifies to not use a tool
elif not request.tool_choice or request.tool_choice == "none":
message = ChatMessage(
role=role, reasoning_content=reasoning_content, content=content
)
message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
# handle when there are tools and tool choice is auto
elif (
@@ -570,16 +507,13 @@ class OmniOpenAIServingChat(OpenAIServingChat):
and self.enable_auto_tools
and self.tool_parser
):
try:
tool_parser = self.tool_parser(tokenizer)
except RuntimeError as e:
logger.exception("Error in tool parser creation.")
return self.create_error_response(str(e))
tool_call_info = tool_parser.extract_tool_calls(
content if content is not None else "", request=request
)
tool_call_info = tool_parser.extract_tool_calls(content if content is not None else "", request=request)
# In the OpenAI API the finish_reason is "tools_called"
# if the tool choice is auto and the model produced a tool
# call. The same is not true for named function calls
@@ -610,37 +544,26 @@ class OmniOpenAIServingChat(OpenAIServingChat):
# undetermined case that is still important to handle
else:
logger.error(
"Error in chat_completion_full_generator - cannot determine"
" if tools should be extracted. Returning a standard chat "
"completion."
)
message = ChatMessage(
role=role, reasoning_content=reasoning_content, content=content
"Error in chat_completion_full_generator - cannot determine if tools should be extracted. "
"Returning a standard chat completion."
)
message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content)
choice_data = ChatCompletionResponseChoice(
index=output.index,
message=message,
logprobs=logprobs,
finish_reason=(
"tool_calls"
if auto_tools_called
else output.finish_reason if output.finish_reason else "stop"
"tool_calls" if auto_tools_called else output.finish_reason if output.finish_reason else "stop"
),
stop_reason=output.stop_reason,
token_ids=(
as_list(output.token_ids) if request.return_token_ids else None
),
token_ids=(as_list(output.token_ids) if request.return_token_ids else None),
)
choices.append(choice_data)
if request.echo:
last_msg_content: Union[str, list[dict[str, str]]] = ""
if (
conversation
and "content" in conversation[-1]
and conversation[-1].get("role") == role
):
if conversation and "content" in conversation[-1] and conversation[-1].get("role") == role:
last_msg_content = conversation[-1]["content"] or ""
if isinstance(last_msg_content, list):
last_msg_content = "\n".join(msg["text"] for msg in last_msg_content)
@@ -653,23 +576,17 @@ class OmniOpenAIServingChat(OpenAIServingChat):
num_prompt_tokens = len(final_res.prompt_token_ids)
if final_res.encoder_prompt_token_ids is not None:
num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
num_generated_tokens = sum(
len(output.token_ids) for output in final_res.outputs
)
num_generated_tokens = sum(len(output.token_ids) for output in final_res.outputs)
usage = UsageInfo(
prompt_tokens=num_prompt_tokens,
completion_tokens=num_generated_tokens,
total_tokens=num_prompt_tokens + num_generated_tokens,
)
if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
usage.prompt_tokens_details = PromptTokenUsageInfo(
cached_tokens=final_res.num_cached_tokens
)
usage.prompt_tokens_details = PromptTokenUsageInfo(cached_tokens=final_res.num_cached_tokens)
prompt_logprobs = clamp_prompt_logprobs(final_res.prompt_logprobs)
prompt_token_ids = (
final_res.prompt_token_ids if request.return_token_ids else None
)
prompt_token_ids = final_res.prompt_token_ids if request.return_token_ids else None
kv_transfer_params = final_res.kv_transfer_params
return choices, usage, prompt_logprobs, prompt_token_ids, kv_transfer_params
@@ -682,8 +599,7 @@ class OmniOpenAIServingChat(OpenAIServingChat):
# Convert numpy array to WAV bytes and encode as base64
if soundfile is None:
raise ImportError(
"soundfile is required for audio generation. "
"Please install it with: pip install soundfile"
"soundfile is required for audio generation. Please install it with: pip install soundfile"
)
# Default sample rate for TTS models (typically 24000 Hz)
vllm_omni/entrypoints/stage_utils.py
+17 -25
View File
@@ -5,7 +5,7 @@ import logging
import os
import pickle
from multiprocessing import shared_memory as _shm
from typing import Any, Dict, Optional, Tuple, Union
from typing import Any
import cloudpickle
from omegaconf import OmegaConf
@@ -13,7 +13,7 @@ from omegaconf import OmegaConf
logger = logging.getLogger(__name__)
def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) -> None:
def set_stage_gpu_devices(stage_id: int, devices: str | int | None) -> None:
"""Configure per-stage CUDA visibility and current device.
Behavior
@@ -26,8 +26,8 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
- Otherwise: set CUDA_VISIBLE_DEVICES to the provided single device string.
"""
try:
selected_physical: Optional[int] = None
logical_idx: Optional[int] = None
selected_physical: int | None = None
logical_idx: int | None = None
if isinstance(devices, str) and "," in devices:
os.environ["CUDA_VISIBLE_DEVICES"] = devices
@@ -42,13 +42,9 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
selected_physical,
)
except Exception as e:
logger.debug(
"[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e
)
logger.debug("[Stage-%s] Failed to parse first CUDA device: %s", stage_id, e)
selected_physical = None
elif isinstance(devices, (int, str)) and (
isinstance(devices, int) or str(devices).isdigit()
):
elif isinstance(devices, (int, str)) and (isinstance(devices, int) or str(devices).isdigit()):
logical_idx = max(0, int(devices))
vis = os.environ.get("CUDA_VISIBLE_DEVICES")
if vis:
@@ -113,9 +109,7 @@ def set_stage_gpu_devices(stage_id: int, devices: Optional[Union[str, int]]) ->
"free": int(free),
}
)
logger.debug(
"[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info
)
logger.debug("[Stage-%s] CUDA devices visible=%s info=%s", stage_id, num, info)
except Exception as e:
logger.debug(
"[Stage-%s] Failed to query CUDA devices: %s",
@@ -132,7 +126,7 @@ def serialize_obj(obj: Any) -> bytes:
return cloudpickle.dumps(obj)
def shm_write_bytes(payload: bytes) -> Dict[str, Any]:
def shm_write_bytes(payload: bytes) -> dict[str, Any]:
"""Write bytes into SharedMemory and return meta dict {name,size}.
Caller should close the segment; the receiver should unlink.
@@ -149,7 +143,7 @@ def shm_write_bytes(payload: bytes) -> Dict[str, Any]:
return meta
def shm_read_bytes(meta: Dict[str, Any]) -> bytes:
def shm_read_bytes(meta: dict[str, Any]) -> bytes:
"""Read bytes from SharedMemory by meta {name,size} and cleanup."""
shm = _shm.SharedMemory(name=meta["name"]) # type: ignore[index]
mv = memoryview(shm.buf)
@@ -176,7 +170,7 @@ def _ensure_parent_dir(path: str) -> None:
pass
def append_jsonl(path: str, record: Dict[str, Any]) -> None:
def append_jsonl(path: str, record: dict[str, Any]) -> None:
"""Append a JSON record as one line to a JSONL file (best-effort).
This is safe to call from multiple processes when each process writes
@@ -193,7 +187,7 @@ def append_jsonl(path: str, record: Dict[str, Any]) -> None:
logger.exception("Failed to append JSONL to %s", path)
def maybe_dump_to_shm(obj: Any, threshold: int) -> Tuple[bool, Any]:
def maybe_dump_to_shm(obj: Any, threshold: int) -> tuple[bool, Any]:
"""Dump object to SHM if serialized size exceeds threshold.
Returns (True, meta) when dumped; otherwise (False, original_obj).
@@ -204,7 +198,7 @@ def maybe_dump_to_shm(obj: Any, threshold: int) -> Tuple[bool, Any]:
return False, obj
def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) -> Any:
def maybe_load_from_ipc(container: dict[str, Any], obj_key: str, shm_key: str) -> Any:
"""Load object from container that may carry SHM or inline object.
Deprecated: prefer `maybe_load_from_ipc_with_metrics` to also obtain
@@ -216,8 +210,8 @@ def maybe_load_from_ipc(container: Dict[str, Any], obj_key: str, shm_key: str) -
def maybe_load_from_ipc_with_metrics(
container: Dict[str, Any], obj_key: str, shm_key: str
) -> tuple[Any, Dict[str, float]]:
container: dict[str, Any], obj_key: str, shm_key: str
) -> tuple[Any, dict[str, float]]:
"""Load object and return (object, metrics) with RX bytes and decode time.
Metrics keys:
@@ -249,15 +243,13 @@ def maybe_load_from_ipc_with_metrics(
}
def encode_for_ipc(
obj: Any, threshold: int, obj_key: str, shm_key: str
) -> Dict[str, Any]:
def encode_for_ipc(obj: Any, threshold: int, obj_key: str, shm_key: str) -> dict[str, Any]:
"""Return a dict payload for IPC: inline (obj_key) or SHM (shm_key).
When serialized size exceeds threshold, returns {shm_key: {name,size}};
otherwise returns {obj_key: obj}.
"""
payload: Dict[str, Any] = {}
payload: dict[str, Any] = {}
use_shm, data = maybe_dump_to_shm(obj, threshold)
if use_shm:
payload[shm_key] = data
@@ -267,7 +259,7 @@ def encode_for_ipc(
# Convert OmegaConf/objects to plain dicts
def _to_dict(x: Any) -> Dict[str, Any]:
def _to_dict(x: Any) -> dict[str, Any]:
try:
if isinstance(x, dict):
return dict(x)
vllm_omni/entrypoints/utils.py
-1
View File
@@ -5,7 +5,6 @@ import os
from pathlib import Path
from omegaconf import OmegaConf
from vllm.transformers_utils.config import get_config
logger = logging.getLogger(__name__)
vllm_omni/inputs/data.py
-1
View File
@@ -7,7 +7,6 @@ except ImportError:
from typing_extensions import NotRequired
import torch
from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt
vllm_omni/inputs/parse.py
+3 -9
View File
@@ -17,15 +17,9 @@ def parse_singleton_prompt_omni(prompt: SingletonPrompt) -> ParsedSingletonPromp
# Priority tokens: When both tokens and embeds exist, keep both and
# follow the tokens path
if "prompt_token_ids" in prompt:
return ParsedTokensPrompt(
type="tokens", content=prompt
) # type: ignore[typeddict-item]
return ParsedTokensPrompt(type="tokens", content=prompt) # type: ignore[typeddict-item]
elif "prompt_embeds" in prompt:
return ParsedEmbedsPrompt(
type="embeds", content=prompt
) # type: ignore[typeddict-item]
return ParsedEmbedsPrompt(type="embeds", content=prompt) # type: ignore[typeddict-item]
elif "prompt" in prompt:
return ParsedTextPrompt(type="text", content=prompt)
raise TypeError(
"inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt"
)
raise TypeError("inputs must be a string, TextPrompt, TokensPrompt, or EmbedsPrompt")
vllm_omni/inputs/preprocess.py
+2 -4
View File
@@ -1,11 +1,11 @@
from typing import Any, Optional, Union
from typing_extensions import assert_never
from vllm.inputs.data import SingletonInputs, SingletonPrompt, TextPrompt, TokensPrompt
from vllm.inputs.preprocess import InputPreprocessor
from vllm.logger import init_logger
from vllm.multimodal.inputs import MultiModalInputs, MultiModalUUIDDict
from vllm_omni.inputs.data import OmniTokenInputs, token_inputs_omni
from vllm_omni.inputs.parse import parse_singleton_prompt_omni
@@ -20,9 +20,7 @@ class OmniInputPreprocessor(InputPreprocessor):
*,
mm_uuids: Optional[MultiModalUUIDDict] = None,
) -> Union[OmniTokenInputs, MultiModalInputs]:
prompt_token_ids = self._truncate_inputs(
parsed_content["prompt_token_ids"], tokenization_kwargs
)
prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs)
prompt_embeds = parsed_content.get("prompt_embeds")
additional_information = parsed_content.get("additional_information")
vllm_omni/model_executor/layers/mrope.py
+43 -178
View File
@@ -4,7 +4,6 @@ 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
@@ -249,9 +248,7 @@ class MRotaryEmbedding(RotaryEmbedding):
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]
):
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
@@ -260,9 +257,7 @@ class MRotaryEmbedding(RotaryEmbedding):
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
)
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],
@@ -275,27 +270,10 @@ class MRotaryEmbedding(RotaryEmbedding):
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
)
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":
@@ -311,36 +289,17 @@ class MRotaryEmbedding(RotaryEmbedding):
)
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
)
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
)
llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)
video_frame_num = 1
else:
@@ -372,9 +331,7 @@ class MRotaryEmbedding(RotaryEmbedding):
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_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()
@@ -424,15 +381,11 @@ class MRotaryEmbedding(RotaryEmbedding):
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
)
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)
torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w)
* video_second_per_grid_t
* tokens_per_second
)
@@ -440,29 +393,15 @@ class MRotaryEmbedding(RotaryEmbedding):
.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
)
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_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()
@@ -512,9 +451,7 @@ class MRotaryEmbedding(RotaryEmbedding):
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
)
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)
@@ -534,21 +471,13 @@ class MRotaryEmbedding(RotaryEmbedding):
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
)
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
):
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
):
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])
@@ -571,58 +500,40 @@ class MRotaryEmbedding(RotaryEmbedding):
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
)
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()
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
)
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
)
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
)
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)
)
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,
@@ -633,38 +544,23 @@ class MRotaryEmbedding(RotaryEmbedding):
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]
)
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
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)
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
)
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]
)
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
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
@@ -673,9 +569,7 @@ class MRotaryEmbedding(RotaryEmbedding):
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)
)
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
@@ -692,25 +586,10 @@ class MRotaryEmbedding(RotaryEmbedding):
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()
)
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()
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)
@@ -731,14 +610,7 @@ class MRotaryEmbedding(RotaryEmbedding):
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)
]
return [list(range(context_len + mrope_position_delta, seq_len + mrope_position_delta)) for _ in range(3)]
@staticmethod
def get_next_input_positions_tensor(
@@ -748,7 +620,6 @@ class MRotaryEmbedding(RotaryEmbedding):
context_len: int,
num_new_tokens: int,
):
values = np.arange(
mrope_position_delta + context_len,
mrope_position_delta + context_len + num_new_tokens,
@@ -779,25 +650,19 @@ class MRotaryEmbedding(RotaryEmbedding):
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
)
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 = (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)
)
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)
vllm_omni/model_executor/model_loader/weight_utils.py
-1
View File
@@ -4,7 +4,6 @@ from typing import Optional, Union
import huggingface_hub
from huggingface_hub import snapshot_download
from vllm.logger import init_logger
from vllm.model_executor.model_loader.weight_utils import DisabledTqdm, get_lock
vllm_omni/model_executor/models/qwen2_5_omni.py
+55 -142
View File
@@ -1,7 +1,8 @@
import glob
import os
from collections.abc import Iterable
from functools import cached_property
from typing import Dict, Iterable, NamedTuple, Optional, Set, Tuple, Union
from typing import NamedTuple, Optional, Union
import numpy as np
import torch
@@ -11,7 +12,6 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
Qwen2_5OmniTalkerConfig,
Qwen2_5OmniThinkerConfig,
)
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP
@@ -29,9 +29,8 @@ from vllm.sequence import IntermediateTensors
from vllm.v1.outputs import SamplerOutput
from vllm.v1.sample.metadata import SamplingMetadata
from vllm.v1.sample.sampler import Sampler
from vllm_omni.model_executor.model_loader.weight_utils import (
download_weights_from_hf_specific,
)
from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific
from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights
TALKER_CODEC_EOS_TOKEN_ID = 8294
@@ -57,7 +56,6 @@ logger = init_logger(__name__)
class Qwen2_5OmniForConditionalGeneration(
nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin
):
def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
super().__init__()
self.have_multimodal_outputs = True
@@ -117,17 +115,15 @@ class Qwen2_5OmniForConditionalGeneration(
architectures=["Qwen2_5OmniToken2WavModel"],
)
# voice resources (loaded on demand)
self._token2wav_conds: Dict[str, torch.Tensor] = {}
self._token2wav_ref_mels: Dict[str, torch.Tensor] = {}
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
(self.thinker.make_empty_intermediate_tensors) if self.model_stage == "thinker" else lambda: None
)
# -------------------- Device utilities --------------------
@@ -176,11 +172,7 @@ class Qwen2_5OmniForConditionalGeneration(
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 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):
@@ -229,9 +221,7 @@ class Qwen2_5OmniForConditionalGeneration(
# 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(
input_ids = torch.zeros(inputs_embeds.shape[1], dtype=torch.long, device=thinker_dev).unsqueeze(
0
) # (1, 0)
added_batch_dim = True
@@ -259,11 +249,7 @@ class Qwen2_5OmniForConditionalGeneration(
# Text-only path
return OmniOutput(
text_hidden_states=(
text_hidden_states.squeeze(0)
if added_batch_dim
else text_hidden_states
),
text_hidden_states=(text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states),
multimodal_outputs=None,
)
@@ -273,9 +259,11 @@ class Qwen2_5OmniForConditionalGeneration(
# Rules:
# - Prefill segments are wrapped with special tokens: [BOS][PAD...][EOS]
# - Decode segments consist of a single non-special token.
# - If additional_information is provided (can be a list split by request or a concatenated tensor plus a list of shapes),
# - If additional_information is provided
# (can be a list split by request or a concatenated tensor plus a list of shapes),
# then for each request, reconstruct the thinker→talker input embeddings for the Prefill segments;
# - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them; otherwise, keep the original embedding.
# - For Decode segments, if per-request auxiliary decode embeddings are provided (optional), add them;
# otherwise, keep the original embedding.
if input_ids is None and additional_information is None:
input_ids = torch.zeros(
@@ -295,9 +283,7 @@ class Qwen2_5OmniForConditionalGeneration(
# ------- Request-scoped additional information (no cross-request concat) -------
request_ids: Optional[list[str]] = kwargs.get("request_ids") # ordered
request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get(
"request_token_spans"
)
request_token_spans: Optional[list[tuple[int, int]]] = kwargs.get("request_token_spans")
addi_by_req: Optional[dict] = kwargs.get("additional_information_by_req_id")
runtime_addi = kwargs.get("runtime_additional_information")
@@ -360,35 +346,24 @@ class Qwen2_5OmniForConditionalGeneration(
req_input_ids, req_embeds = self._thinker_to_talker_prefill(
voice_type=voice_type,
output_prompt_embeds=tr.to(inputs_embeds.dtype).to(
self._module_device(self.model)
),
output_prompt_embeds=tr.to(inputs_embeds.dtype).to(self._module_device(self.model)),
output_token_ids=otoks,
thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(
self._module_device(self.model)
),
thinker_prompt_embeds=pe.to(inputs_embeds.dtype).to(self._module_device(self.model)),
prompt_token_ids=ptoks,
)
seg_len = min(span_len, req_embeds.shape[0])
inputs_embeds[s : s + seg_len] = req_embeds[:seg_len]
if (
isinstance(req_input_ids, torch.Tensor)
and req_input_ids.numel() == seg_len
):
if isinstance(req_input_ids, torch.Tensor) and req_input_ids.numel() == seg_len:
input_ids[s : s + seg_len] = req_input_ids
# Prepare per-request reply queue for subsequent decode: drop first row
if tr.ndim == 2 and tr.shape[0] > 0:
update_by_req_id.setdefault(rid, {})[
"thinker_reply_part_per_request"
] = (tr[1:].detach().to("cpu").contiguous())
update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = (
tr[1:].detach().to("cpu").contiguous()
)
# ------- Decode: span_len == 1 -------
if (
not is_profile
and isinstance(request_ids, list)
and isinstance(request_token_spans, list)
):
if not is_profile and isinstance(request_ids, list) and isinstance(request_token_spans, list):
for idx_req, rid in enumerate(request_ids):
s, e = request_token_spans[idx_req]
if (int(e) - int(s)) != 1:
@@ -398,27 +373,15 @@ class Qwen2_5OmniForConditionalGeneration(
# A) runtime queue
q = None
if isinstance(rid, str):
q = runtime_addi_by_req.get(rid, {}).get(
"thinker_reply_part_per_request"
)
q = runtime_addi_by_req.get(rid, {}).get("thinker_reply_part_per_request")
if isinstance(q, torch.Tensor) and q.numel() > 0:
step_vec = q[0:1]
new_q = q[1:].detach().to("cpu").contiguous()
update_by_req_id.setdefault(rid, {})[
"thinker_reply_part_per_request"
] = new_q
update_by_req_id.setdefault(rid, {})["thinker_reply_part_per_request"] = new_q
else:
# B) per-request provided decode vector (optional)
info = (
addi_by_req.get(rid, {})
if isinstance(addi_by_req, dict)
else {}
)
dv = (
info.get("decode_output_prompt_embeds")
if isinstance(info, dict)
else None
)
info = addi_by_req.get(rid, {}) if isinstance(addi_by_req, dict) else {}
dv = info.get("decode_output_prompt_embeds") if isinstance(info, dict) else None
if isinstance(dv, torch.Tensor) and dv.numel() > 0:
step_vec = dv[0:1] if dv.ndim == 2 else dv.view(1, -1)
elif (
@@ -433,9 +396,7 @@ class Qwen2_5OmniForConditionalGeneration(
if isinstance(step_vec, torch.Tensor) and step_vec.numel() > 0:
one_id = input_ids[s : s + 1]
_, one_embed = self._thinker_to_talker_decode_one_step(
output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(
self._module_device(self.model)
),
output_prompt_embeds=step_vec.to(inputs_embeds.dtype).to(self._module_device(self.model)),
output_token_ids=one_id,
)
inputs_embeds[s] = one_embed[0]
@@ -449,15 +410,12 @@ class Qwen2_5OmniForConditionalGeneration(
multimodal_outputs: dict = None
# Return updates if any
if update_by_req_id:
multimodal_outputs = {
"additional_information_update_by_req_id": update_by_req_id
}
multimodal_outputs = {"additional_information_update_by_req_id": update_by_req_id}
if sampling_metadata is not None:
# the padding token id is set to text model's pad token id, which do not match with the talker model's word embedding size
sampling_metadata.prompt_token_ids[
sampling_metadata.prompt_token_ids == 152064
] = 8448
# the padding token id is set to text model's pad token id,
# which do not match with the talker model's word embedding size
sampling_metadata.prompt_token_ids[sampling_metadata.prompt_token_ids == 152064] = 8448
return OmniOutput(
text_hidden_states=talker_hidden,
@@ -479,9 +437,7 @@ class Qwen2_5OmniForConditionalGeneration(
code = code[1:] if code[0] == TALKER_CODEC_BOS_TOKEN_ID else code
audio_tensor = self.generate_audio(code, voice_type)
return OmniOutput(
text_hidden_states=None, multimodal_outputs={"audio": audio_tensor}
)
return OmniOutput(text_hidden_states=None, multimodal_outputs={"audio": audio_tensor})
return OmniOutput(
text_hidden_states=torch.cat(
@@ -492,9 +448,7 @@ class Qwen2_5OmniForConditionalGeneration(
).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]
)
torch.tensor([TALKER_CODEC_BOS_TOKEN_ID, TALKER_CODEC_EOS_TOKEN_ID])
.to(torch.bfloat16)
.to(self._module_device(self.model))
)
@@ -621,7 +575,6 @@ class Qwen2_5OmniForConditionalGeneration(
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
@@ -644,9 +597,7 @@ class Qwen2_5OmniForConditionalGeneration(
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) + [
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,
]
@@ -657,9 +608,7 @@ class Qwen2_5OmniForConditionalGeneration(
][-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))
torch.Tensor(prompt_token_ids_processed).to(torch.int64).to(self._module_device(self.talker))
)
return prompt_token_ids_processed, prompt_embeds
@@ -668,14 +617,12 @@ class Qwen2_5OmniForConditionalGeneration(
output_prompt_embeds,
output_token_ids,
):
processed_output_token_embeds = (
output_prompt_embeds + self.talker.get_input_embeddings(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]
) -> Optional[torch.Tensor]:
def compute_logits(self, hidden_states: Union[torch.Tensor, OmniOutput]) -> Optional[torch.Tensor]:
# Handle OmniOutput type
if isinstance(hidden_states, OmniOutput):
hidden_states = hidden_states.text_hidden_states
@@ -691,9 +638,7 @@ class Qwen2_5OmniForConditionalGeneration(
# Use thinker model for sampling
return self.model.sample(logits, sampling_metadata)
def generate_speech(
self, text_tokens: torch.Tensor, voice_type: str = "default"
) -> torch.Tensor:
def generate_speech(self, text_tokens: torch.Tensor, voice_type: str = "default") -> torch.Tensor:
"""
Generate speech from text tokens using the talker and token2wav models.
This method is kept for backward compatibility and direct speech generation.
@@ -706,9 +651,7 @@ class Qwen2_5OmniForConditionalGeneration(
Audio tensor
"""
# Generate codec tokens using talker model
talker_output = self.talker(
input_ids=None, positions=None, inputs_embeds=text_tokens
)
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)
@@ -735,9 +678,7 @@ class Qwen2_5OmniForConditionalGeneration(
# 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"
):
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
@@ -756,12 +697,8 @@ class Qwen2_5OmniForConditionalGeneration(
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()
}
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 = hf_model_folder
if isinstance(model_path, str) and os.path.isdir(model_path):
@@ -773,24 +710,14 @@ class Qwen2_5OmniForConditionalGeneration(
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"))
):
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"))
):
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
)
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]:
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:
@@ -824,9 +751,7 @@ class Qwen2_5OmniForConditionalGeneration(
if codec.ndim == 1:
codec = codec.unsqueeze(0)
else:
codec = torch.as_tensor(
codec_tokens, dtype=torch.long, device=token2wav_dev
).unsqueeze(0)
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)
@@ -841,9 +766,7 @@ class Qwen2_5OmniForConditionalGeneration(
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
)
y_all = torch.randn((1, total_mel, mel_dim), dtype=ref_mel.dtype, device=token2wav_dev)
logger.info(
"Currently, we do not use the chunked process, we only use the "
@@ -855,9 +778,7 @@ class Qwen2_5OmniForConditionalGeneration(
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:
if (chunk_code_length > 0 and chunk_code_length % chunk_size == 0) or finished:
chunk_ends.append(i)
# Number of chunks in mel domain
@@ -886,7 +807,7 @@ class Qwen2_5OmniForConditionalGeneration(
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]:
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 = []
@@ -915,20 +836,14 @@ class Qwen2_5OmniForConditionalGeneration(
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"
]
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)
)
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")
@@ -949,9 +864,7 @@ class Qwen2_5OmniForConditionalGeneration(
allow_patterns=["*.pt"],
)
self._init_token2wav_model(hf_model_folder)
t2w_loaded = self.token2wav.load_weights(
token2wav_weights, os.path.join(hf_model_folder, "spk_dict.pt")
)
t2w_loaded = self.token2wav.load_weights(token2wav_weights, os.path.join(hf_model_folder, "spk_dict.pt"))
t2w_loaded = add_prefix_to_loaded_weights(t2w_loaded, "token2wav")
loaded_weights.update(t2w_loaded)
vllm_omni/model_executor/models/qwen2_5_omni_talker.py
+15 -45
View File
@@ -1,24 +1,17 @@
from collections.abc import Iterable
from functools import cached_property
from typing import Iterable, Optional, Set, Tuple, Union
from typing import Optional, 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 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 # unused import
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.layers.linear import ColumnParallelLinear
from vllm.model_executor.models.interfaces import (
MultiModalEmbeddings,
SupportsMultiModal,
SupportsPP,
)
from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
from vllm.model_executor.models.qwen2_5_omni_thinker import (
Qwen2_5OmniConditionalGenerationMixin,
Qwen2_5OmniThinkerDummyInputsBuilder,
@@ -72,9 +65,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
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
)
vllm_config.model_config.hf_text_config = vllm_config.model_config.hf_config.talker_config
else:
self.config = config
@@ -92,9 +83,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
hf_config=getattr(self.config, "text_config", self.config),
architectures=["Qwen2ForCausalLM_old"],
)
self.make_empty_intermediate_tensors = (
self.language_model.make_empty_intermediate_tensors
)
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)
@@ -122,7 +111,6 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
) -> 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(
@@ -145,9 +133,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
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"
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() # unused variable
if intermediate_tensors is not None:
@@ -184,7 +170,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
) -> Optional[SamplerOutput]:
return self.language_model.sample(logits, sampling_metadata)
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = AutoWeightsLoader(
self,
skip_prefixes=["thinker.", "token2wav."],
@@ -220,31 +206,15 @@ class Qwen2_5OmniTalkerForConditionalGeneration(
# 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
)
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 []
vllm_omni/model_executor/models/qwen2_5_omni_thinker.py
+54 -180
View File
@@ -6,25 +6,13 @@ 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.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.interfaces import (
MultiModalEmbeddings,
SupportsMultiModal,
SupportsPP,
)
from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
from vllm.model_executor.models.qwen2_5_vl import (
Qwen2_5_VisionTransformer,
Qwen2_5_VLImageEmbeddingInputs,
@@ -64,15 +52,11 @@ from vllm.multimodal.parse import (
MultiModalDataItems,
MultiModalDataParser,
)
from vllm.multimodal.processing import (
BaseMultiModalProcessor,
PlaceholderFeaturesInfo,
PromptReplacement,
PromptUpdate,
)
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_omni.model_executor.layers.mrope import MRotaryEmbedding
try:
@@ -93,17 +77,13 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
video_grid_sizes = video_grid_thw.prod(-1)
return dict(
input_audio_features=MultiModalFieldConfig.flat_from_sizes(
"audio", audio_feature_lengths, dim=1
),
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
),
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"),
@@ -111,7 +91,6 @@ def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
def _parse_audio_data(
self,
data: Union[dict[str, torch.Tensor], ModalityData[ImageItem]],
@@ -127,10 +106,7 @@ class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
return super()._parse_audio_data(data)
class Qwen2_5OmniThinkerProcessingInfo(
Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo
):
class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo):
def get_hf_config(self):
return self.ctx.get_hf_config(Qwen2_5OmniConfig).thinker_config
@@ -151,10 +127,7 @@ class Qwen2_5OmniThinkerProcessingInfo(
return {"audio": None, "image": None, "video": None}
class Qwen2_5OmniThinkerDummyInputsBuilder(
BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]
):
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)
@@ -166,11 +139,7 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
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
)
return audio_token * num_audios + image_token * num_images + video_token * num_videos
def get_dummy_mm_data(
self,
@@ -191,17 +160,11 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
* 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
)
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
),
"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,
@@ -213,15 +176,10 @@ class Qwen2_5OmniThinkerDummyInputsBuilder(
return mm_data
class Qwen2_5OmniThinkerMultiModalProcessor(
BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]
):
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
)
return Qwen2_5OmniThinkerMultiModalDataParser(target_sr=feature_extractor.sampling_rate)
def _call_hf_processor(
self,
@@ -252,14 +210,9 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
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)
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
):
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
@@ -299,9 +252,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
prompt_ids,
mm_item_counts,
)
self._validate_mm_placeholders(
mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
)
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)
@@ -315,9 +266,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
mm_prompt_updates,
mm_item_counts,
)
self._validate_mm_placeholders(
mm_placeholders, mm_item_counts, use_audio_in_video=use_audio_in_video
)
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)
@@ -350,15 +299,11 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
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_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_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1))
audio_output_lengths = audio_output_lens.tolist()
# number of audios read from video.
@@ -372,8 +317,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
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"
f"The audio {audio} (len={len(audio)}) is too short to be represented inside the model"
)
return [audio_token_id] * num_features
@@ -392,9 +336,7 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
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
]
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
@@ -509,39 +451,27 @@ class Qwen2_5OmniThinkerMultiModalProcessor(
class Qwen2_5OmniConditionalGenerationMixin:
def _validate_and_reshape_mm_tensor(
self, mm_input: object, name: str, dim: int = 0
) -> torch.Tensor:
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)}"
)
raise ValueError(f"Incorrect type of {name}. 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]:
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
)
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)}"
)
raise ValueError(f"Incorrect type of audio input features. Got type: {type(input_audio_features)}")
return Qwen2AudioInputs(
input_features=input_audio_features,
audio_feature_lengths=audio_feature_lengths,
@@ -560,18 +490,11 @@ class Qwen2_5OmniConditionalGenerationMixin:
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"
)
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)}"
)
raise ValueError(f"Incorrect type of image pixel values. Got type: {type(pixel_values)}")
return Qwen2_5_VLImagePixelInputs(
type="pixel_values",
@@ -580,18 +503,11 @@ class Qwen2_5OmniConditionalGenerationMixin:
)
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"
)
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)}"
)
raise ValueError(f"Incorrect type of image embeddings. Got type: {type(image_embeds)}")
return Qwen2_5_VLImageEmbeddingInputs(
type="image_embeds",
image_embeds=image_embeds,
@@ -610,12 +526,8 @@ class Qwen2_5OmniConditionalGenerationMixin:
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"
)
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",
@@ -624,18 +536,11 @@ class Qwen2_5OmniConditionalGenerationMixin:
)
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"
)
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)}"
)
raise ValueError(f"Incorrect type of video embeddings. Got type: {type(video_embeds)}")
return Qwen2_5_VLVideoEmbeddingInputs(
type="video_embeds",
video_embeds=video_embeds,
@@ -648,24 +553,20 @@ class Qwen2_5OmniConditionalGenerationMixin:
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
)
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_feat_lengths, audio_output_lengths = self.audio_tower._get_feat_extract_output_lengths(
audio_feature_lengths
)
audio_outputs = self.audio_tower(
@@ -676,9 +577,7 @@ class Qwen2_5OmniConditionalGenerationMixin:
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, ...]:
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)
@@ -758,8 +657,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
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."
"exactly same result as the transformers implementation in the audio tower part."
)
self.audio_tower = Qwen2_5OmniAudioEncoder(thinker_config.audio_config)
@@ -777,9 +675,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
architectures=["Qwen2ForCausalLM"],
)
self.make_empty_intermediate_tensors = (
self.language_model.make_empty_intermediate_tensors
)
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 = {}
@@ -787,34 +683,18 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
# 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
)
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 []
@@ -845,7 +725,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
) -> 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(
@@ -897,9 +776,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
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
)
inputs_embeds = merge_multimodal_embeddings(input_ids, inputs_embeds, embeddings, placeholder_token_id)
return inputs_embeds
def forward(
@@ -921,10 +798,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
text_inputs_embeds = self.get_input_embeddings(
input_ids,
(
[
(torch.zeros_like(embeddings), modality)
for embeddings, modality in multimodal_embeddings
]
[(torch.zeros_like(embeddings), modality) for embeddings, modality in multimodal_embeddings]
if multimodal_embeddings is not None
else None
),
vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py
+88 -244
View File
@@ -3,7 +3,8 @@
############################
import math
from typing import Iterable, List, Optional, Set, Tuple, Union
from collections.abc import Iterable
from typing import Optional, Union
import numpy as np
import torch
@@ -16,24 +17,17 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import (
Qwen2_5OmniDiTConfig,
Qwen2_5OmniToken2WavConfig,
)
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import (
Qwen2_5OmniPreTrainedModel,
)
from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniPreTrainedModel
# Bring in HF base classes, configs and utilities used below
from transformers.utils.logging import get_logger as _hf_get_logger
from vllm.config import VllmConfig
from vllm.logger import init_logger
from vllm.model_executor.models.interfaces import SupportsPP
from vllm.model_executor.models.utils import (
AutoWeightsLoader as _Vllm_AutoWeightsLoader,
)
from vllm.model_executor.models.utils import AutoWeightsLoader as _Vllm_AutoWeightsLoader
from vllm.model_executor.models.utils import WeightsMapper as _Vllm_WeightsMapper
from vllm.model_executor.models.utils import (
init_vllm_registered_model as _Vllm_init_vllm_registered_model,
)
from vllm.model_executor.models.utils import maybe_prefix as _Vllm_maybe_prefix
from vllm.model_executor.models.utils import init_vllm_registered_model as _vllm_init_vllm_registered_model
from vllm.model_executor.models.utils import maybe_prefix as _vllm_maybe_prefix
from vllm.sequence import IntermediateTensors
from vllm.v1.outputs import SamplerOutput
from vllm.v1.sample.metadata import SamplingMetadata
@@ -218,9 +212,7 @@ class AttentiveStatisticsPooling(nn.Module):
def _compute_statistics(self, x, m, dim=2):
mean = (m * x).sum(dim)
std = torch.sqrt(
(m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps)
)
std = torch.sqrt((m * (x - mean.unsqueeze(dim)).pow(2)).sum(dim).clamp(self.eps))
return mean, std
def forward(self, hidden_states):
@@ -282,9 +274,7 @@ class SqueezeExcitationRes2NetBlock(nn.Module):
kernel_size=1,
dilation=1,
)
self.res2net_block = Res2NetBlock(
out_channels, out_channels, res2net_scale, kernel_size, dilation
)
self.res2net_block = Res2NetBlock(out_channels, out_channels, res2net_scale, kernel_size, dilation)
self.tdnn2 = TimeDelayNetBlock(
out_channels,
out_channels,
@@ -312,13 +302,10 @@ class ECAPA_TimeDelayNet(torch.nn.Module):
def __init__(self, config: Qwen2_5OmniDiTConfig):
super().__init__()
if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(
config.enc_channels
) != len(config.enc_dilations):
raise ValueError(
"enc_channels, enc_kernel_sizes and enc_dilations should have "
"same length"
)
if len(config.enc_channels) != len(config.enc_kernel_sizes) or len(config.enc_channels) != len(
config.enc_dilations
):
raise ValueError("enc_channels, enc_kernel_sizes and enc_dilations should have same length")
self.channels = config.enc_channels
self.blocks = nn.ModuleList()
@@ -412,26 +399,14 @@ class DiTInputEmbedding(nn.Module):
):
if apply_cfg:
hidden_states = torch.cat([hidden_states, hidden_states], dim=0)
speaker_embedding = torch.cat(
[speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0
)
condition_vector = torch.cat(
[condition_vector, torch.zeros_like(condition_vector)], dim=0
)
speaker_embedding = torch.cat([speaker_embedding, torch.zeros_like(speaker_embedding)], dim=0)
condition_vector = torch.cat([condition_vector, torch.zeros_like(condition_vector)], dim=0)
code_embed = torch.cat([code_embed, code_embed_uncond], dim=0)
elif drop_audio_cond: # cfg for cond audio
condition_vector = torch.zeros_like(condition_vector)
speaker_embedding = torch.zeros_like(speaker_embedding)
condition_vector = (
self.spk_encoder(condition_vector)
.unsqueeze(1)
.repeat(1, hidden_states.size(1), 1)
)
hidden_states = self.proj(
torch.cat(
(hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1
)
)
condition_vector = self.spk_encoder(condition_vector).unsqueeze(1).repeat(1, hidden_states.size(1), 1)
hidden_states = self.proj(torch.cat((hidden_states, condition_vector, code_embed, speaker_embedding), dim=-1))
return hidden_states
@@ -465,13 +440,9 @@ class Qwen2_5_OmniAdaLayerNormZero(nn.Module):
def forward(self, hidden_states, emb=None):
emb = self.linear(self.silu(emb))
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(
emb, 6, dim=1
)
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = torch.chunk(emb, 6, dim=1)
hidden_states = (
self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
)
hidden_states = self.norm(hidden_states) * (1 + scale_msa[:, None]) + shift_msa[:, None]
return hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp
@@ -490,9 +461,7 @@ class Qwen2_5_OmniAdaLayerNormZero_Final(nn.Module):
emb = self.linear(self.silu(emb))
scale, shift = torch.chunk(emb, 2, dim=1)
hidden_states = (
self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
)
hidden_states = self.norm(hidden_states) * (1 + scale)[:, None, :] + shift[:, None, :]
return hidden_states
@@ -572,9 +541,7 @@ class DiTAttention(nn.Module):
self.to_k = nn.Linear(config.hidden_size, self.inner_dim)
self.to_v = nn.Linear(config.hidden_size, self.inner_dim)
self.to_out = nn.ModuleList(
[nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)]
)
self.to_out = nn.ModuleList([nn.Linear(self.inner_dim, config.hidden_size), nn.Dropout(config.dropout)])
def forward(
self,
@@ -600,9 +567,7 @@ class DiTAttention(nn.Module):
# Due to training process, only first head is applied with RoPE,
# will be fixed at next release
cos, sin = position_embeddings
query[:, :1], key[:, :1] = apply_rotary_pos_emb(
query[:, :1], key[:, :1], cos, sin
)
query[:, :1], key[:, :1] = apply_rotary_pos_emb(query[:, :1], key[:, :1], cos, sin)
attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]
attention_weights, _ = attention_interface(
@@ -616,9 +581,7 @@ class DiTAttention(nn.Module):
# mask. e.g. inference got a batch with different target durations,
# mask out the padding
attention_weights = attention_weights.reshape(
batch_size, -1, self.heads * head_dim
)
attention_weights = attention_weights.reshape(batch_size, -1, self.heads * head_dim)
attention_weights = attention_weights.to(query.dtype)
# linear proj
@@ -648,9 +611,7 @@ class DiTTimestepEmbedding(nn.Module):
def __init__(self, dim, freq_embed_dim=256):
super().__init__()
self.time_embed = SinusPositionEmbedding(freq_embed_dim)
self.time_mlp = nn.ModuleList(
[nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)]
)
self.time_mlp = nn.ModuleList([nn.Linear(freq_embed_dim, dim), nn.SiLU(), nn.Linear(dim, dim)])
def forward(self, timestep): # noqa: F821
time_hidden = self.time_embed(timestep)
@@ -661,29 +622,21 @@ class DiTTimestepEmbedding(nn.Module):
class DiTDecoderLayer(nn.Module):
def __init__(
self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0
):
def __init__(self, config: Qwen2_5OmniDiTConfig, look_ahead_block=0, look_backward_block=0):
super().__init__()
self.attn_norm = Qwen2_5_OmniAdaLayerNormZero(config.hidden_size)
self.attn = DiTAttention(config)
self.look_ahead_block = look_ahead_block
self.look_backward_block = look_backward_block
self.ff_norm = nn.LayerNorm(
config.hidden_size, elementwise_affine=False, eps=1e-6
)
self.ff = DiTMLP(
dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout
)
self.ff_norm = nn.LayerNorm(config.hidden_size, elementwise_affine=False, eps=1e-6)
self.ff = DiTMLP(dim=config.hidden_size, mult=config.ff_mult, dropout=config.dropout)
def forward(
self, hidden_states, timestep, position_embeddings=None, block_diff=None
): # x: noised input, t: time embedding
# pre-norm & modulation for attention input
norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(
hidden_states, emb=timestep
)
norm, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.attn_norm(hidden_states, emb=timestep)
# attention
attn_output = self.attn(
@@ -696,9 +649,7 @@ class DiTDecoderLayer(nn.Module):
# process attention output for input x
hidden_states = hidden_states + gate_msa.unsqueeze(1) * attn_output
norm = (
self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
)
norm = self.ff_norm(hidden_states) * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
ff_output = self.ff(norm)
hidden_states = hidden_states + gate_mlp.unsqueeze(1) * ff_output
@@ -741,16 +692,14 @@ class SnakeBeta(nn.Module):
beta = self.beta.unsqueeze(0).unsqueeze(-1)
alpha = torch.exp(alpha)
beta = torch.exp(beta)
hidden_states = hidden_states + (
1.0 / (beta + self.no_div_by_zero)
) * torch.pow(torch.sin(hidden_states * alpha), 2)
hidden_states = hidden_states + (1.0 / (beta + self.no_div_by_zero)) * torch.pow(
torch.sin(hidden_states * alpha), 2
)
return hidden_states
def kaiser_sinc_filter1d(
cutoff: float, half_width: float, kernel_size: int
) -> torch.Tensor:
def kaiser_sinc_filter1d(cutoff: float, half_width: float, kernel_size: int) -> torch.Tensor:
"""Generates a 1D Kaiser-windowed sinc filter.
Args:
@@ -775,9 +724,7 @@ def kaiser_sinc_filter1d(
else:
beta = 0.0
kaiser_window = torch.kaiser_window(
kernel_size, beta=beta, periodic=False, dtype=torch.float32
)
kaiser_window = torch.kaiser_window(kernel_size, beta=beta, periodic=False, dtype=torch.float32)
# Compute time indices
if is_even:
@@ -787,9 +734,7 @@ def kaiser_sinc_filter1d(
# Compute sinc filter
if cutoff == 0:
return torch.zeros(
(1, 1, kernel_size), dtype=torch.float32
) # Ensures correct shape
return torch.zeros((1, 1, kernel_size), dtype=torch.float32) # Ensures correct shape
sinc_filter = torch.sinc(2 * cutoff * time_indices)
normalized_filter = 2 * cutoff * kaiser_window * sinc_filter
@@ -804,27 +749,19 @@ class UpSample1d(nn.Module):
def __init__(self, ratio=2, kernel_size=None):
super().__init__()
self.ratio = ratio
self.kernel_size = (
int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
)
self.kernel_size = int(6 * ratio // 2) * 2 if kernel_size is None else kernel_size
self.stride = ratio
self.pad = self.kernel_size // ratio - 1
self.pad_left = self.pad * self.stride + (self.kernel_size - self.stride) // 2
self.pad_right = (
self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
)
self.pad_right = self.pad * self.stride + (self.kernel_size - self.stride + 1) // 2
filter = kaiser_sinc_filter1d(
cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size
)
filter = kaiser_sinc_filter1d(cutoff=0.5 / ratio, half_width=0.6 / ratio, kernel_size=self.kernel_size)
self.register_buffer("filter", filter, persistent=False)
def forward(self, hidden_states):
channels = hidden_states.shape[1]
hidden_states_dtype = hidden_states.dtype
hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(
self.filter.dtype
)
hidden_states = F.pad(hidden_states, (self.pad, self.pad), mode="replicate").to(self.filter.dtype)
hidden_states = self.ratio * F.conv_transpose1d(
hidden_states,
self.filter.expand(channels, -1, -1),
@@ -857,9 +794,7 @@ class DownSample1d(nn.Module):
def forward(self, hidden_states):
channels = hidden_states.shape[1]
hidden_states_dtype = hidden_states.dtype
hidden_states = F.pad(
hidden_states, (self.pad_left, self.pad_right), mode="replicate"
).to(self.filter.dtype)
hidden_states = F.pad(hidden_states, (self.pad_left, self.pad_right), mode="replicate").to(self.filter.dtype)
out = F.conv1d(
hidden_states,
self.filter.expand(channels, -1, -1),
@@ -960,15 +895,10 @@ class AMPBlock(torch.nn.Module):
]
)
self.num_layers = len(self.convs1) + len(
self.convs2
) # total number of conv layers
self.num_layers = len(self.convs1) + len(self.convs2) # total number of conv layers
self.activations = nn.ModuleList(
[
TorchActivation1d(activation=SnakeBeta(channels))
for _ in range(self.num_layers)
]
[TorchActivation1d(activation=SnakeBeta(channels)) for _ in range(self.num_layers)]
)
def _get_padding(self, kernel_size, dilation=1):
@@ -1001,9 +931,7 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
self.num_residual_blocks = len(config.resblock_kernel_sizes)
self.num_upsample_layers = len(config.upsample_rates)
self.conv_pre = nn.Conv1d(
config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3
)
self.conv_pre = nn.Conv1d(config.mel_dim, config.upsample_initial_channel, 7, 1, padding=3)
# Removing extra ModuleList breaks official state dict
ups = [
@@ -1018,9 +946,7 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
)
]
)
for layer_idx, (stride, kernel_size) in enumerate(
zip(config.upsample_rates, config.upsample_kernel_sizes)
)
for layer_idx, (stride, kernel_size) in enumerate(zip(config.upsample_rates, config.upsample_kernel_sizes))
]
self.ups = nn.ModuleList(ups)
@@ -1032,16 +958,12 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
dilation,
)
for layer_idx in range(self.num_upsample_layers)
for kernel_size, dilation in zip(
config.resblock_kernel_sizes, config.resblock_dilation_sizes
)
for kernel_size, dilation in zip(config.resblock_kernel_sizes, config.resblock_dilation_sizes)
]
)
self.activation_post = TorchActivation1d(
activation=SnakeBeta(
config.upsample_initial_channel // (2**self.num_upsample_layers)
)
activation=SnakeBeta(config.upsample_initial_channel // (2**self.num_upsample_layers))
)
self.conv_post = nn.Conv1d(
config.upsample_initial_channel // (2**self.num_upsample_layers),
@@ -1081,9 +1003,7 @@ class Qwen2_5OmniToken2WavBigVGANModel(Qwen2_5OmniPreTrainedModel):
for layer_index in range(self.num_upsample_layers):
hidden_representation = self.ups[layer_index][0](hidden_representation)
residual_output = sum(
self.resblocks[layer_index * self.num_residual_blocks + block_index](
hidden_representation
)
self.resblocks[layer_index * self.num_residual_blocks + block_index](hidden_representation)
for block_index in range(self.num_residual_blocks)
)
residual_output = residual_output / self.num_residual_blocks
@@ -1111,11 +1031,7 @@ class RungeKutta4ODESolver:
value_start,
function_value_start=None,
):
k1 = (
function_value_start
if function_value_start is not None
else function(time_start, value_start)
)
k1 = function_value_start if function_value_start is not None else function(time_start, value_start)
k2 = function(
time_start + time_step * self._one_third,
value_start + time_step * k1 * self._one_third,
@@ -1141,9 +1057,7 @@ class RungeKutta4ODESolver:
function_value_start,
)
def _linear_interpolation(
self, time_start, time_end, value_start, value_end, time_point
):
def _linear_interpolation(self, time_start, time_end, value_start, value_end, time_point):
if time_point == time_start:
return value_start
if time_point == time_end:
@@ -1164,15 +1078,10 @@ class RungeKutta4ODESolver:
current_value = self.initial_value
for time_start, time_end in zip(time_points[:-1], time_points[1:]):
time_step = time_end - time_start
delta_value, _ = self._compute_step(
self.function, time_start, time_step, time_end, current_value
)
delta_value, _ = self._compute_step(self.function, time_start, time_step, time_end, current_value)
next_value = current_value + delta_value
while (
current_index < len(time_points)
and time_end >= time_points[current_index]
):
while current_index < len(time_points) and time_end >= time_points[current_index]:
solution[current_index] = self._linear_interpolation(
time_start,
time_end,
@@ -1203,9 +1112,7 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
self.repeats = config.repeats
self.time_embed = DiTTimestepEmbedding(config.hidden_size)
self.text_embed = DiTCodecEmbedding(
config.num_embeds, config.emb_dim, config.repeats
)
self.text_embed = DiTCodecEmbedding(config.num_embeds, config.emb_dim, config.repeats)
self.input_embed = DiTInputEmbedding(config)
self.rotary_embed = Qwen2_5OmniDiTRotaryEmbedding(config.head_dim)
@@ -1225,16 +1132,12 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
)
)
self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(
config.hidden_size
) # final modulation
self.norm_out = Qwen2_5_OmniAdaLayerNormZero_Final(config.hidden_size) # final modulation
self.proj_out = nn.Linear(config.hidden_size, config.mel_dim)
def _create_block_diff(self, hidden_states):
batch, seq_len = hidden_states.shape[0], hidden_states.shape[1]
block_indices = (
torch.arange(seq_len, device=hidden_states.device) // self.block_size
) # [seq_length]
block_indices = torch.arange(seq_len, device=hidden_states.device) // self.block_size # [seq_length]
block_i = block_indices.unsqueeze(1) # [seq_length, 1]
block_j = block_indices.unsqueeze(0) # [1, seq_length]
@@ -1259,12 +1162,8 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
# Compute embeddings
time_embedding = self.time_embed(time_step)
text_embedding = self.text_embed(
quantized_code, drop_code=False if apply_cfg else drop_code
)
text_embedding_unconditioned = (
self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
)
text_embedding = self.text_embed(quantized_code, drop_code=False if apply_cfg else drop_code)
text_embedding_unconditioned = self.text_embed(quantized_code, drop_code=True) if apply_cfg else None
hidden_states = self.input_embed(
hidden_states,
@@ -1304,17 +1203,11 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
guidance_scale=0.5,
sway_coefficient=-1.0,
):
noise_initialization = torch.randn(
[1, 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype
)
noise_initialization = torch.randn([1, 30000, self.mel_dim], dtype=reference_mel_spectrogram.dtype)
maximum_duration = quantized_code.shape[1] * self.repeats
initial_state = noise_initialization[:, :maximum_duration].to(
quantized_code.device
)
initial_state = noise_initialization[:, :maximum_duration].to(quantized_code.device)
batch_size = reference_mel_spectrogram.shape[0]
conditioning_vector = conditioning_vector.unsqueeze(1).repeat(
1, maximum_duration, 1
)
conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, maximum_duration, 1)
if batch_size != 1:
raise ValueError("Only batch size = 1 is currently supported")
@@ -1341,10 +1234,7 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
apply_cfg=True,
)
guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
return (
guided_prediction
+ (guided_prediction - null_prediction) * guidance_scale
)
return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
initial_time = 0
time_embedding = torch.linspace(
@@ -1356,13 +1246,9 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
)
if sway_coefficient is not None:
time_embedding += sway_coefficient * (
torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding
)
time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
ode_solver = RungeKutta4ODESolver(
function=ode_function, initial_value=initial_state
)
ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
solution_trajectory = ode_solver.integrate(time_embedding)
generated_waveform = solution_trajectory[-1]
@@ -1393,9 +1279,7 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
"""
initial_state = y0.to(quantized_code.device)
batch_size = reference_mel_spectrogram.shape[0]
conditioning_vector = conditioning_vector.unsqueeze(1).repeat(
1, initial_state.shape[1], 1
)
conditioning_vector = conditioning_vector.unsqueeze(1).repeat(1, initial_state.shape[1], 1)
if batch_size != 1:
raise ValueError("Only batch size = 1 is currently supported")
@@ -1422,10 +1306,7 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
apply_cfg=True,
)
guided_prediction, null_prediction = torch.chunk(model_output, 2, dim=0)
return (
guided_prediction
+ (guided_prediction - null_prediction) * guidance_scale
)
return guided_prediction + (guided_prediction - null_prediction) * guidance_scale
initial_time = 0
time_embedding = torch.linspace(
@@ -1437,13 +1318,9 @@ class Qwen2_5OmniToken2WavDiTModel(Qwen2_5OmniPreTrainedModel):
)
if sway_coefficient is not None:
time_embedding += sway_coefficient * (
torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding
)
time_embedding += sway_coefficient * (torch.cos(torch.pi / 2 * time_embedding) - 1 + time_embedding)
ode_solver = RungeKutta4ODESolver(
function=ode_function, initial_value=initial_state
)
ode_solver = RungeKutta4ODESolver(function=ode_function, initial_value=initial_state)
solution_trajectory = ode_solver.integrate(time_embedding)
generated_waveform = solution_trajectory[-1]
@@ -1479,8 +1356,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
attn_impl = "sdpa"
elif config._attn_implementation == "eager":
logger.warning_once(
"Qwen2_5OmniToken2WavModel does not support eager attention "
"implementation, fall back to sdpa"
"Qwen2_5OmniToken2WavModel does not support eager attention implementation, fall back to sdpa"
)
attn_impl = "sdpa"
self.code2wav_dit_model = Qwen2_5OmniToken2WavDiTModel._from_config(
@@ -1503,9 +1379,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
# codec embedding size for masking EOS out-of-range
try:
self.codec_embed_size = (
self.code2wav_dit_model.text_embed.codec_embed.weight.size(0)
)
self.codec_embed_size = self.code2wav_dit_model.text_embed.codec_embed.weight.size(0)
except Exception:
self.codec_embed_size = -1
@@ -1587,7 +1461,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
steps: int,
prev_generated: torch.Tensor,
finished: bool = False,
) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
) -> tuple[Optional[torch.Tensor], torch.Tensor]:
"""Streaming per small chunk: returns (mel_or_None, audio_slice)."""
start_index = max(i * self.chunk_size - self.past_cache_size, 0)
end_index = min(
@@ -1595,16 +1469,8 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
codec_all.shape[1] * self.factor,
)
y0 = (
y_all[:, start_index:end_index]
.reshape(1, -1, self.code2wav_dit_model.mel_dim)
.contiguous()
)
codec = (
codec_all[:, start_index // self.factor : end_index // self.factor]
.reshape(1, -1)
.contiguous()
)
y0 = y_all[:, start_index:end_index].reshape(1, -1, self.code2wav_dit_model.mel_dim).contiguous()
codec = codec_all[:, start_index // self.factor : end_index // self.factor].reshape(1, -1).contiguous()
# generate mel for current window (B, mel_dim, T)
generated = self.process_chunk_dit_batch(
@@ -1634,9 +1500,9 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
y_all: torch.Tensor,
i: int,
steps: int,
prev_generated: Union[torch.Tensor, List[torch.Tensor]],
prev_generated: Union[torch.Tensor, list[torch.Tensor]],
finished: bool = False,
) -> Tuple[Union[torch.Tensor, List[torch.Tensor]], torch.Tensor]:
) -> tuple[Union[torch.Tensor, list[torch.Tensor]], torch.Tensor]:
"""High-level chunk API aligning to qwen2_code2wav_dit signature."""
if not isinstance(prev_generated, torch.Tensor):
prev_generated = prev_generated[0] if len(prev_generated) > 0 else None
@@ -1661,7 +1527,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
finished: bool,
prev_generated: Optional[torch.Tensor],
generated: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor]:
) -> tuple[torch.Tensor, torch.Tensor]:
"""
Align mel and audio boundaries for 50Hz-like streaming.
@@ -1678,21 +1544,13 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
mel = generated[:, :, : self.chunk_size]
elif finished:
mel_trim = generated[:, :, self.past_cache_size :]
mel = torch.cat(
[prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2
)
mel = torch.cat([prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2)
else:
if start_index == 0:
mel_trim = generated[
:, :, i * self.chunk_size : -self.future_cache_size
]
mel_trim = generated[:, :, i * self.chunk_size : -self.future_cache_size]
else:
mel_trim = generated[
:, :, self.past_cache_size : -self.future_cache_size
]
mel = torch.cat(
[prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2
)
mel_trim = generated[:, :, self.past_cache_size : -self.future_cache_size]
mel = torch.cat([prev_generated[:, :, -self.future_size * 2 :], mel_trim], dim=2)
audio = self.code2wav_bigvgan_model(mel)
if i == 0:
@@ -1700,11 +1558,7 @@ class Qwen2_5OmniToken2WavModel(Qwen2_5OmniPreTrainedModel):
elif finished:
audio_output = audio[self.future_size * self.vocoder_hop :]
else:
audio_output = audio[
self.future_size
* self.vocoder_hop : -self.future_size
* self.vocoder_hop
]
audio_output = audio[self.future_size * self.vocoder_hop : -self.future_size * self.vocoder_hop]
return mel, audio_output
@@ -1728,9 +1582,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
self.config = vllm_config.model_config.hf_config
# Initialize underlying HF Token2Wav model via registry
self.token2wav = _Vllm_init_vllm_registered_model(
self.token2wav = _vllm_init_vllm_registered_model(
vllm_config=vllm_config,
prefix=_Vllm_maybe_prefix(prefix, "token2wav_model"),
prefix=_vllm_maybe_prefix(prefix, "token2wav_model"),
hf_config=self.config,
architectures=["Qwen2_5OmniToken2WavDiTModel"],
)
@@ -1782,9 +1636,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
) -> Optional[SamplerOutput]:
return None
def load_weights_without_buffers(
self, weights: Iterable[Tuple[str, torch.Tensor]]
) -> Set[str]:
def load_weights_without_buffers(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
loader = _Vllm_AutoWeightsLoader(self)
loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
# Log load summary
@@ -1822,9 +1674,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
return registers
# remove buffers from the weights and reload them after loading weights
def remove_buffers_from_weights(
self, weights: Iterable[Tuple[str, torch.Tensor]], buffers: dict
):
def remove_buffers_from_weights(self, weights: Iterable[tuple[str, torch.Tensor]], buffers: dict):
weights_to_load = []
for key, value in weights:
if key in buffers:
@@ -1844,9 +1694,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
loaded_buffers.add(name)
return loaded_buffers
def load_weights(
self, weights: Iterable[Tuple[str, torch.Tensor]], spk_dict_path: str
) -> Set[str]:
def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]], spk_dict_path: str) -> set[str]:
buffers = self.find_all_registers()
weights_to_load = self.remove_buffers_from_weights(weights, buffers)
loaded = self.load_weights_without_buffers(weights_to_load)
@@ -1874,9 +1722,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
)
@torch.inference_mode()
def process_chunk_bigvgan_batch(
self, mel_batch: torch.Tensor
) -> Optional[torch.Tensor]:
def process_chunk_bigvgan_batch(self, mel_batch: torch.Tensor) -> Optional[torch.Tensor]:
# BigVGAN is not part of this wrapper; return None for parity.
return None
@@ -1891,7 +1737,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
steps: int,
prev_generated: torch.Tensor,
finished: bool = False,
) -> Tuple[Optional[torch.Tensor], torch.Tensor]:
) -> tuple[Optional[torch.Tensor], torch.Tensor]:
mel = self.token2wav(
code=codec_all,
conditioning=conditioning,
@@ -1909,9 +1755,9 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
y_all: torch.Tensor,
i: int,
steps: int,
prev_generated: Union[torch.Tensor, List[torch.Tensor]],
prev_generated: Union[torch.Tensor, list[torch.Tensor]],
finished: bool = False,
) -> Tuple[Union[torch.Tensor, List[torch.Tensor]], torch.Tensor]:
) -> tuple[Union[torch.Tensor, list[torch.Tensor]], torch.Tensor]:
_mel, out = self.process_little_chunk(
conditioning=conditioning,
reference_mel=reference_mel,
@@ -1919,9 +1765,7 @@ class Qwen2_5OmniToken2WavForConditionalGenerationVLLM(nn.Module, SupportsPP):
y_all=y_all,
i=i,
steps=steps,
prev_generated=(
prev_generated if isinstance(prev_generated, torch.Tensor) else None
),
prev_generated=(prev_generated if isinstance(prev_generated, torch.Tensor) else None),
finished=finished,
)
return _mel if _mel is not None else prev_generated, out
vllm_omni/model_executor/models/qwen2_old.py
+26 -68
View File
@@ -1,9 +1,9 @@
from typing import Iterable, Optional, Set, Tuple, Union
from collections.abc import Iterable
from typing import Optional, 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
@@ -11,23 +11,13 @@ 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.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.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.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.model_executor.models.interfaces import SupportsLoRA, SupportsPP
from vllm.model_executor.models.utils import (
AutoWeightsLoader,
@@ -48,7 +38,6 @@ logger = init_logger(__name__)
class Qwen2MLP(nn.Module):
def __init__(
self,
hidden_size: int,
@@ -73,10 +62,7 @@ class Qwen2MLP(nn.Module):
prefix=f"{prefix}.down_proj",
)
if hidden_act != "silu":
raise ValueError(
f"Unsupported activation: {hidden_act}. "
"Only silu is supported for now."
)
raise ValueError(f"Unsupported activation: {hidden_act}. Only silu is supported for now.")
self.act_fn = SiluAndMul()
def forward(self, x):
@@ -87,7 +73,6 @@ class Qwen2MLP(nn.Module):
class Qwen2Attention(nn.Module):
def __init__(
self,
hidden_size: int,
@@ -98,7 +83,7 @@ class Qwen2Attention(nn.Module):
head_dim: Optional[int] = None,
cache_config: Optional[CacheConfig] = None,
quant_config: Optional[QuantizationConfig] = None,
rope_scaling: Optional[Tuple] = None,
rope_scaling: Optional[tuple] = None,
prefix: str = "",
attn_type: str = AttentionType.DECODER,
) -> None:
@@ -173,7 +158,6 @@ class Qwen2Attention(nn.Module):
class Qwen2DecoderLayer(nn.Module):
def __init__(
self,
config: Qwen2Config,
@@ -217,16 +201,14 @@ class Qwen2DecoderLayer(nn.Module):
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
)
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]:
) -> tuple[torch.Tensor, torch.Tensor]:
# Self Attention
if residual is None:
residual = hidden_states
@@ -255,7 +237,6 @@ class Qwen2DecoderLayer(nn.Module):
}
)
class Qwen2Model(nn.Module):
def __init__(
self,
*,
@@ -270,27 +251,20 @@ class Qwen2Model(nn.Module):
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"
):
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,
)
f"but `max_window_layers` = {config.max_window_layers} is less than "
f"`num_hidden_layers` = {config.num_hidden_layers}. Please open an issue "
"to discuss this feature."
)
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
):
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),
@@ -348,13 +322,11 @@ class Qwen2Model(nn.Module):
residual,
)
if not get_pp_group().is_last_rank:
return IntermediateTensors(
{"hidden_states": hidden_states, "residual": residual}
)
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]:
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"),
@@ -364,19 +336,15 @@ class Qwen2Model(nn.Module):
("gate_up_proj", "up_proj", 1),
]
params_dict = dict(self.named_parameters(remove_duplicate=False))
loaded_params: Set[str] = set()
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)
):
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]
)
loaded_weight = loaded_weight if loaded_weight.dim() == 0 else loaded_weight[0]
weight_loader(param, loaded_weight)
loaded_params.add(scale_name)
continue
@@ -433,9 +401,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
self.lora_config = lora_config
self.quant_config = quant_config
self.model = Qwen2Model(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model"))
if get_pp_group().is_last_rank:
if config.tie_word_embeddings:
@@ -453,9 +419,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
self.logits_processor = LogitsProcessor(config.vocab_size)
self.sampler = Sampler()
self.make_empty_intermediate_tensors = (
self.model.make_empty_intermediate_tensors
)
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)
@@ -467,9 +431,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
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
)
hidden_states = self.model(input_ids, positions, intermediate_tensors, inputs_embeds)
return hidden_states
def compute_logits(
@@ -487,7 +449,7 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP):
next_tokens = self.sampler(logits, sampling_metadata)
return next_tokens
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]) -> Set[str]:
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),
@@ -521,9 +483,7 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
self.lora_config = lora_config
self.quant_config = quant_config
self.model = Qwen2Model(
vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
)
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
@@ -554,9 +514,7 @@ class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP):
) -> Optional[PoolerOutput]:
return self._pooler(hidden_states, pooling_metadata)
def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):
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.")
)
weights = ((name, data) for name, data in weights if not name.startswith("lm_head."))
self.model.load_weights(weights)
vllm_omni/model_executor/models/registry.py
+1 -5
View File
@@ -1,8 +1,4 @@
from vllm.model_executor.models.registry import (
_VLLM_MODELS,
_LazyRegisteredModel,
_ModelRegistry,
)
from vllm.model_executor.models.registry import _VLLM_MODELS, _LazyRegisteredModel, _ModelRegistry
_OMNI_MODELS = {
"Qwen2_5OmniForConditionalGeneration": (
vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py
+10 -23
View File
@@ -1,8 +1,8 @@
from typing import Union
import torch
from vllm.inputs import TextPrompt
from vllm_omni.inputs.data import OmniTokensPrompt
TALKER_CODEC_PAD_TOKEN_ID = 8292
@@ -10,9 +10,7 @@ TALKER_CODEC_START_TOKEN_ID = 8293
TALKER_CODEC_END_TOKEN_ID = 8294
def thinker2talker(
stage_list, engine_input_source, prompt: Union[OmniTokensPrompt, TextPrompt] = None
):
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]
@@ -25,8 +23,7 @@ def thinker2talker(
if not isinstance(prompt, list):
prompt = [prompt]
multi_modal_data = {
thinker_output.request_id: p.get("multi_modal_data", None)
for thinker_output, p in zip(thinker_outputs, prompt)
thinker_output.request_id: p.get("multi_modal_data", None) for thinker_output, p in zip(thinker_outputs, prompt)
}
for i, thinker_output in enumerate(thinker_outputs):
@@ -36,32 +33,22 @@ def thinker2talker(
prompt_token_ids_len = len(prompt_token_ids)
latent = output.multimodal_output["latent"]
thinker_hidden_states = latent.clone().detach().to(latent.device)
addtional_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
),
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,
"thinker_result_shape": list(
thinker_hidden_states[prompt_token_ids_len:].shape
),
"prompt_embeds_shape": list(
thinker_hidden_states[:prompt_token_ids_len].shape
),
"thinker_result_shape": list(thinker_hidden_states[prompt_token_ids_len:].shape),
"prompt_embeds_shape": list(thinker_hidden_states[:prompt_token_ids_len].shape),
}
talker_inputs.append(
OmniTokensPrompt(
prompt_token_ids=[TALKER_CODEC_START_TOKEN_ID]
+ [TALKER_CODEC_PAD_TOKEN_ID] * (len(prompt_token_ids))
+ [TALKER_CODEC_END_TOKEN_ID],
additional_information=addtional_information,
additional_information=additional_information,
multi_modal_data=(
multi_modal_data[thinker_output.request_id]
if multi_modal_data is not None
else None
multi_modal_data[thinker_output.request_id] if multi_modal_data is not None else None
),
mm_processor_kwargs=None,
)
vllm_omni/outputs.py
-1
View File
@@ -2,7 +2,6 @@ from dataclasses import dataclass
from typing import Optional
import torch
from vllm.outputs import RequestOutput
from vllm.v1.outputs import ModelRunnerOutput
vllm_omni/patch.py
+4 -11
View File
@@ -1,11 +1,10 @@
import sys
from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt
from vllm.model_executor.layers.rotary_embedding import (
MRotaryEmbedding as _OriginalMRotaryEmbedding,
)
from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding as _OriginalMRotaryEmbedding
from vllm.v1.engine import EngineCoreRequest as _OriginalEngineCoreRequest
from vllm.v1.request import Request as _OriginalRequest
from vllm_omni.engine import OmniEngineCoreRequest
from vllm_omni.inputs.data import OmniTokensPrompt
from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding
@@ -14,15 +13,9 @@ from vllm_omni.request import OmniRequest
for module_name, module in sys.modules.items():
if hasattr(module, "TokensPrompt") and module.TokensPrompt == _OriginalTokensPrompt:
module.TokensPrompt = OmniTokensPrompt
if (
hasattr(module, "MRotaryEmbedding")
and module.MRotaryEmbedding == _OriginalMRotaryEmbedding
):
if hasattr(module, "MRotaryEmbedding") and module.MRotaryEmbedding == _OriginalMRotaryEmbedding:
module.MRotaryEmbedding = MRotaryEmbedding
if hasattr(module, "Request") and module.Request == _OriginalRequest:
module.Request = OmniRequest
if (
hasattr(module, "EngineCoreRequest")
and module.EngineCoreRequest == _OriginalEngineCoreRequest
):
if hasattr(module, "EngineCoreRequest") and module.EngineCoreRequest == _OriginalEngineCoreRequest:
module.EngineCoreRequest = OmniEngineCoreRequest
vllm_omni/request.py
+3 -12
View File
@@ -6,11 +6,7 @@ from vllm.v1.structured_output.request import StructuredOutputRequest
if TYPE_CHECKING:
from vllm.v1.core.kv_cache_utils import BlockHash
from vllm_omni.engine import (
AdditionalInformationPayload,
OmniEngineCoreRequest,
PromptEmbedsPayload,
)
from vllm_omni.engine import AdditionalInformationPayload, OmniEngineCoreRequest, PromptEmbedsPayload
class OmniRequest(Request):
@@ -25,9 +21,7 @@ class OmniRequest(Request):
# Serialized prompt embeddings payload (optional)
self.prompt_embeds: Optional[PromptEmbedsPayload] = prompt_embeds
# Serialized additional information payload (optional)
self.additional_information: Optional[AdditionalInformationPayload] = (
additional_information
)
self.additional_information: Optional[AdditionalInformationPayload] = additional_information
@classmethod
def from_engine_core_request(
@@ -35,7 +29,6 @@ class OmniRequest(Request):
request: OmniEngineCoreRequest,
block_hasher: Optional[Callable[["Request"], list["BlockHash"]]],
) -> "Request":
return cls(
request_id=request.request_id,
client_index=request.client_index,
@@ -47,9 +40,7 @@ class OmniRequest(Request):
arrival_time=request.arrival_time,
lora_request=request.lora_request,
structured_output_request=(
StructuredOutputRequest(sampling_params=request.sampling_params)
if request.sampling_params
else None
StructuredOutputRequest(sampling_params=request.sampling_params) if request.sampling_params else None
),
cache_salt=request.cache_salt,
priority=request.priority,
vllm_omni/worker/gpu_ar_model_runner.py
+53 -127
View File
@@ -6,11 +6,10 @@ and also outputs sampled tokens.
from __future__ import annotations
from typing import Any, Optional, Union
from typing import Any
import numpy as np
import torch
from vllm.forward_context import BatchDescriptor
from vllm.logger import init_logger
from vllm.v1.core.sched.output import SchedulerOutput
@@ -28,6 +27,7 @@ from vllm.v1.worker.gpu_model_runner import (
)
from vllm.v1.worker.ubatch_utils import UBatchSlices
from vllm.v1.worker.utils import is_residual_scattered_for_sp
from vllm_omni.engine import AdditionalInformationPayload, PromptEmbedsPayload
from vllm_omni.outputs import OmniModelRunnerOutput
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
@@ -51,23 +51,22 @@ class GPUARModelRunner(OmniGPUModelRunner):
def _preprocess(
self,
scheduler_output: "SchedulerOutput",
scheduler_output: SchedulerOutput,
num_scheduled_tokens_np: np.ndarray,
intermediate_tensors: Optional[IntermediateTensors] = None,
ubatch_slices: Optional[UBatchSlices] = None,
num_tokens_after_padding: Optional[torch.Tensor] = None,
intermediate_tensors: IntermediateTensors | None = None,
ubatch_slices: UBatchSlices | None = None,
num_tokens_after_padding: torch.Tensor | None = None,
) -> tuple[
int,
int,
Optional[torch.Tensor],
Optional[torch.Tensor],
Optional[torch.Tensor],
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor,
Optional[IntermediateTensors],
IntermediateTensors | None,
dict[str, Any],
Optional[dict[str, dict]],
dict[str, dict] | None,
]:
num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens
if ubatch_slices:
assert num_tokens_after_padding is not None
@@ -80,12 +79,8 @@ class GPUARModelRunner(OmniGPUModelRunner):
# _prepare_inputs may reorder the batch, so we must gather multi
# modal outputs after that to ensure the correct order
per_req_additional_information: Optional[dict[str, dict]] = None
if (
self.supports_mm_inputs
and get_pp_group().is_first_rank
and not self.model_config.is_encoder_decoder
):
per_req_additional_information: dict[str, dict] | None = None
if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
# Run the multimodal encoder if any.
self._execute_mm_encoder(scheduler_output)
mm_embeds = self._gather_mm_embeddings(scheduler_output)
@@ -113,9 +108,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
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]
)
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])
@@ -123,26 +116,22 @@ class GPUARModelRunner(OmniGPUModelRunner):
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
src = pe_cpu[num_computed_tokens : num_computed_tokens + overlay_len].to(
dtype=self.dtype, device=self.device, non_blocking=True
)
start_offset = int(self.query_start_loc.cpu[req_index])
self.inputs_embeds[start_offset : start_offset + overlay_len].copy_(src)
# Build per-request additional information (no cross-request concat)
if addi_cpu is not None and isinstance(addi_cpu, dict):
req_info: dict[str, object] = {}
for k, v in addi_cpu.items():
if isinstance(v, torch.Tensor):
# For prefill tokens, pass only the scheduled slice; for decode or no scheduled tokens, pass whole tensor
# For prefill tokens, pass only the scheduled slice;
# for decode or no scheduled tokens, pass whole tensor
if overlay_len > 0:
try:
seg = (
v[
num_computed_tokens : num_computed_tokens
+ overlay_len
]
v[num_computed_tokens : num_computed_tokens + overlay_len]
.detach()
.to("cpu")
.contiguous()
@@ -177,11 +166,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
# If a batch only has token ids, then including the embedding layer
# in the CUDA graph will be more performant (like in the else case
# below).
token_ids_idx = (
self.is_token_ids.gpu[:num_scheduled_tokens]
.nonzero(as_tuple=False)
.squeeze(1)
)
token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1)
# Some tokens ids may need to become embeds
if token_ids_idx.numel() > 0:
token_ids = self.input_ids.gpu[token_ids_idx]
@@ -211,10 +196,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
num_input_tokens, intermediate_tensors, True
)
if (
self.model_config.is_encoder_decoder
and scheduler_output.scheduled_encoder_inputs
):
if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
encoder_inputs = self._extract_encoder_inputs(scheduler_output)
model_kwargs.update(encoder_inputs)
@@ -233,9 +215,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
@torch.inference_mode()
def execute_model(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[OmniModelRunnerOutput, AsyncModelRunnerOutput, IntermediateTensors]:
scheduler_output: SchedulerOutput,
intermediate_tensors: IntermediateTensors | None = None,
) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors:
with record_function_or_nullcontext("Preprocess"):
with self.synchronize_input_prep():
# Update persistent batch states.
@@ -247,9 +229,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
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
)
req_id = getattr(nr, "req_id", None) or getattr(nr, "request_id", None)
if req_id is None:
continue
# prompt_embeds
@@ -258,9 +238,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
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")
)
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)
@@ -279,19 +257,11 @@ class GPUARModelRunner(OmniGPUModelRunner):
if isinstance(payload_info, dict):
# Already decoded
info_dict = payload_info
elif isinstance(
payload_info, AdditionalInformationPayload
):
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
)
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:
@@ -303,18 +273,14 @@ class GPUARModelRunner(OmniGPUModelRunner):
info_dict,
)
except Exception as e:
logger.error(
f"Error decoding prompt_embeds / additional_information: {e}"
)
logger.error(f"Error decoding prompt_embeds / additional_information: {e}")
pass
if not scheduler_output.total_num_scheduled_tokens:
if not has_kv_transfer_group():
# Return empty ModelRunnerOutput if no work to do.
return EMPTY_MODEL_RUNNER_OUTPUT
return self.kv_connector_no_forward(
scheduler_output, self.vllm_config
)
return self.kv_connector_no_forward(scheduler_output, self.vllm_config)
if self.cache_config.kv_sharing_fast_prefill:
assert not self.input_batch.num_prompt_logprobs, (
"--kv-sharing-fast-prefill produces incorrect "
@@ -355,12 +321,8 @@ class GPUARModelRunner(OmniGPUModelRunner):
uniform_decode = (max_query_len == self.uniform_decode_query_len) and (
num_scheduled_tokens == self.input_batch.num_reqs * max_query_len
)
batch_descriptor = BatchDescriptor(
num_tokens=num_input_tokens, uniform_decode=uniform_decode
)
cudagraph_runtime_mode, batch_descriptor = (
self.cudagraph_dispatcher.dispatch(batch_descriptor)
)
batch_descriptor = BatchDescriptor(num_tokens=num_input_tokens, uniform_decode=uniform_decode)
cudagraph_runtime_mode, batch_descriptor = self.cudagraph_dispatcher.dispatch(batch_descriptor)
# This is currently to get around the assert in the DPMetadata
# where it wants `num_tokens_across_dp` to align with `num_tokens`
@@ -382,29 +344,21 @@ class GPUARModelRunner(OmniGPUModelRunner):
record_function_or_nullcontext("Forward"),
self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
):
model_kwargs_extra = {}
# Pass per-request additional information map for this step (no concat)
if per_req_additional_information:
model_kwargs_extra["additional_information_by_req_id"] = (
per_req_additional_information
)
model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information
# Always pass per-request runtime additional_information (persisted in request state)
try:
per_req_runtime_info = []
for req_id in self.input_batch.req_ids:
req_state = self.requests.get(req_id)
info = (
getattr(req_state, "additional_information_cpu", None)
if req_state is not None
else None
)
info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None
per_req_runtime_info.append(info if isinstance(info, dict) else {})
model_kwargs_extra["runtime_additional_information"] = (
per_req_runtime_info
)
model_kwargs_extra["runtime_additional_information"] = per_req_runtime_info
model_kwargs_extra["request_ids"] = self.input_batch.req_ids
# Pass each request's token span within the flattened sequence for this step, enabling the model to map decode/prefill by request
# Pass each request's token span within the flattened sequence for this step,
# enabling the model to map decode/prefill by request
req_token_spans = []
for req_index in range(len(self.input_batch.req_ids)):
start_offset = int(self.query_start_loc.cpu[req_index])
@@ -434,9 +388,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
hidden_states = model_output
aux_hidden_states = None
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(
hidden_states
)
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
# The model side may return per-request additional_information updates (model-agnostic channel).
# Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order;
# the runner merges it into the corresponding request's additional_information_cpu for subsequent decode.
@@ -446,21 +398,15 @@ class GPUARModelRunner(OmniGPUModelRunner):
or "additional_information_update_by_req_id" in multimodal_outputs
):
# Option A: list[dict] in batch order
updates_list = multimodal_outputs.get(
"additional_information_update"
)
updates_list = multimodal_outputs.get("additional_information_update")
if isinstance(updates_list, list):
for idx, upd in enumerate(updates_list):
if not isinstance(upd, dict) or idx >= len(
self.input_batch.req_ids
):
if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids):
continue
req_id = self.input_batch.req_ids[idx]
self._merge_additional_information_update(req_id, upd)
# Option B: dict[str, dict] keyed by req_id
updates_map = multimodal_outputs.get(
"additional_information_update_by_req_id"
)
updates_map = multimodal_outputs.get("additional_information_update_by_req_id")
if isinstance(updates_map, dict):
for req_id, upd in updates_map.items():
if not isinstance(upd, dict):
@@ -474,7 +420,6 @@ class GPUARModelRunner(OmniGPUModelRunner):
information update: {e}, with the multimodal_outputs as {multimodal_outputs}"
)
if not self.broadcast_pp_output:
# Common case.
if not get_pp_group().is_last_rank:
# Return the intermediate tensors.
@@ -484,9 +429,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
if self.is_pooling_model:
# Return the pooling output.
output = self._pool(
hidden_states, num_scheduled_tokens, num_scheduled_tokens_np
)
output = self._pool(hidden_states, num_scheduled_tokens, num_scheduled_tokens_np)
output.kv_connector_output = kv_connector_output
return output
@@ -498,9 +441,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
if not get_pp_group().is_last_rank:
all_gather_tensors = {
"residual": not is_residual_scattered_for_sp(
self.vllm_config, num_input_tokens
)
"residual": not is_residual_scattered_for_sp(self.vllm_config, num_input_tokens)
}
get_pp_group().send_tensor_dict(
hidden_states.tensors,
@@ -524,9 +465,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
# Apply structured output bitmasks if present
if scheduler_output.grammar_bitmask is not None:
apply_grammar_bitmask(
scheduler_output, self.input_batch, logits, self.device
)
apply_grammar_bitmask(scheduler_output, self.input_batch, logits, self.device)
with record_function_or_nullcontext("Sample"):
sampler_output = self._sample(logits, spec_decode_metadata)
@@ -558,12 +497,9 @@ class GPUARModelRunner(OmniGPUModelRunner):
and self.speculative_config.draft_model_config is not None
and self.speculative_config.draft_model_config.max_model_len is not None
):
effective_drafter_max_model_len = (
self.speculative_config.draft_model_config.max_model_len
)
effective_drafter_max_model_len = self.speculative_config.draft_model_config.max_model_len
input_fits_in_drafter = spec_decode_common_attn_metadata and (
spec_decode_common_attn_metadata.seq_lens.max()
+ self.speculative_config.num_speculative_tokens
spec_decode_common_attn_metadata.seq_lens.max() + self.speculative_config.num_speculative_tokens
<= effective_drafter_max_model_len
)
if use_padded_batch_for_eagle and input_fits_in_drafter:
@@ -588,11 +524,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
num_scheduled_tokens,
)
if (
self.speculative_config
and not use_padded_batch_for_eagle
and input_fits_in_drafter
):
if self.speculative_config and not use_padded_batch_for_eagle and input_fits_in_drafter:
# ngram and other speculative decoding methods use the sampled
# tokens on the CPU, so they are run after bookkeeping.
propose_draft_token_ids(valid_sampled_token_ids)
@@ -602,12 +534,10 @@ class GPUARModelRunner(OmniGPUModelRunner):
# Convert to per-request tensors on CPU
hidden_states_cpu = hidden_states.detach().to("cpu").contiguous()
pooler_output: list[Optional[torch.Tensor]] = []
pooler_output: list[torch.Tensor | None] = []
prev_logits_index = 0
for logits_index in logits_indices:
pooler_output.append(
hidden_states_cpu[prev_logits_index : logits_index + 1]
)
pooler_output.append(hidden_states_cpu[prev_logits_index : logits_index + 1])
prev_logits_index = logits_index + 1
output = OmniModelRunnerOutput(
@@ -616,11 +546,7 @@ class GPUARModelRunner(OmniGPUModelRunner):
sampled_token_ids=valid_sampled_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
),
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,
)
vllm_omni/worker/gpu_ar_worker.py
+3 -8
View File
@@ -2,13 +2,13 @@ import gc
import os
import torch
from vllm.model_executor import set_random_seed
from vllm.platforms import current_platform
from vllm.utils import GiB_bytes, MemorySnapshot
from vllm.v1.utils import report_usage_stats
from vllm.v1.worker.gpu_worker import Worker as GPUWorker
from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner
@@ -43,10 +43,7 @@ class GPUARWorker(GPUWorker):
# take current memory snapshot
self.init_snapshot = MemorySnapshot()
self.requested_memory = (
self.init_snapshot.total_memory
* self.cache_config.gpu_memory_utilization
)
self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization
if self.init_snapshot.free_memory < self.requested_memory:
def GiB(b):
@@ -65,9 +62,7 @@ class GPUARWorker(GPUWorker):
raise RuntimeError(f"Not support device type: {self.device_config.device}")
# Construct the model runner
self.model_runner: GPUARModelRunner = GPUARModelRunner(
self.vllm_config, self.device
)
self.model_runner: GPUARModelRunner = GPUARModelRunner(self.vllm_config, self.device)
if self.rank == 0:
# If usage stat is enabled, collect relevant info.
vllm_omni/worker/gpu_diffusion_model_runner.py
+46 -93
View File
@@ -2,11 +2,10 @@ from __future__ import annotations
import gc
import logging
from typing import Any, List, Optional, Union
from typing import Any
import numpy as np
import torch
from vllm.config import CUDAGraphMode
from vllm.forward_context import BatchDescriptor
from vllm.multimodal.inputs import MultiModalKwargs
@@ -24,6 +23,7 @@ from vllm.v1.worker.gpu_model_runner import (
)
from vllm.v1.worker.ubatch_utils import UBatchSlices
from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
from vllm_omni.outputs import OmniModelRunnerOutput
from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner
@@ -40,21 +40,20 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
def _preprocess(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
ubatch_slices: Optional[UBatchSlices] = None,
num_tokens_after_padding: Optional[torch.Tensor] = None,
scheduler_output: SchedulerOutput,
intermediate_tensors: IntermediateTensors | None = None,
ubatch_slices: UBatchSlices | None = None,
num_tokens_after_padding: torch.Tensor | None = None,
) -> tuple[
int,
int,
Optional[torch.Tensor],
Optional[torch.Tensor],
Optional[torch.Tensor],
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor | None,
torch.Tensor,
Optional[IntermediateTensors],
IntermediateTensors | None,
dict[str, Any],
]:
# Input token count for this iteration (not used by diffusion, but
# retained to keep DP padding/ordering consistent)
num_input_tokens = scheduler_output.total_num_scheduled_tokens
@@ -63,11 +62,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
# _prepare_inputs may reorder the batch, so we must gather multi
# modal outputs after that to ensure the correct order
if (
self.supports_mm_inputs
and get_pp_group().is_first_rank
and not self.model_config.is_encoder_decoder
):
if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder:
# Run the multimodal encoder if any.
self._execute_mm_encoder(scheduler_output)
mm_embeds = self._gather_mm_embeddings(scheduler_output)
@@ -102,11 +97,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
# If a batch only has token ids, then including the embedding layer
# in the CUDA graph will be more performant (like in the else case
# below).
token_ids_idx = (
self.is_token_ids.gpu[:num_input_tokens]
.nonzero(as_tuple=False)
.squeeze(1)
)
token_ids_idx = self.is_token_ids.gpu[:num_input_tokens].nonzero(as_tuple=False).squeeze(1)
# Some tokens ids may need to become embeds
if token_ids_idx.numel() > 0:
token_ids = self.input_ids.gpu[token_ids_idx]
@@ -136,10 +127,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
num_input_tokens, intermediate_tensors, True
)
if (
self.model_config.is_encoder_decoder
and scheduler_output.scheduled_encoder_inputs
):
if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs:
encoder_inputs = self._extract_encoder_inputs(scheduler_output)
model_kwargs.update(encoder_inputs)
@@ -157,9 +145,9 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
@torch.inference_mode()
def execute_model(
self,
scheduler_output: "SchedulerOutput",
intermediate_tensors: Optional[IntermediateTensors] = None,
) -> Union[OmniModelRunnerOutput, IntermediateTensors]:
scheduler_output: SchedulerOutput,
intermediate_tensors: IntermediateTensors | None = None,
) -> OmniModelRunnerOutput | IntermediateTensors:
with record_function_or_nullcontext("Preprocess"):
with self.synchronize_input_prep():
# Update persistent batch states.
@@ -213,7 +201,6 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
record_function_or_nullcontext("Forward"),
self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output,
):
outputs = self._run_diffusion(
input_ids=input_ids,
positions=positions,
@@ -225,24 +212,18 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
_, multimodal_outputs = self.extract_multimodal_outputs(outputs)
# Ensure one tensor per request, map to CPU for output struct
pooler_output: List[Optional[torch.Tensor]] = []
pooler_output: list[torch.Tensor | None] = []
if isinstance(multimodal_outputs, torch.Tensor):
# If model returned a single stacked tensor, split by requests
assert multimodal_outputs.shape[0] == self.input_batch.num_reqs
for i in range(self.input_batch.num_reqs):
pooler_output.append(
multimodal_outputs[i].detach().to("cpu").contiguous()
)
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
)
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
)
pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None)
else:
raise RuntimeError("Unsupported diffusion output type")
@@ -272,11 +253,11 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
*,
input_ids: torch.Tensor,
positions: torch.Tensor,
intermediate_tensors: Optional[IntermediateTensors],
inputs_embeds: Optional[torch.Tensor],
intermediate_tensors: IntermediateTensors | None,
inputs_embeds: torch.Tensor | None,
multimodal_kwargs: dict,
logits_indices: torch.Tensor,
) -> Union[torch.Tensor, list[torch.Tensor]]:
) -> torch.Tensor | list[torch.Tensor]:
"""Runs the diffusion process and returns per-request tensors.
Tries model interfaces in the following order for maximal compatibility:
@@ -302,15 +283,15 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
# 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."
"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,
cudagraph_runtime_mode: Optional[CUDAGraphMode] = None,
cudagraph_runtime_mode: CUDAGraphMode | None = None,
force_attention: bool = False,
uniform_decode: bool = False,
skip_eplb: bool = False,
@@ -389,7 +370,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
num_tokens_across_dp = num_tokens_after_padding
num_tokens_after_padding = int(num_tokens_after_padding[0].item())
attn_metadata: Optional[PerLayerAttnMetadata] = None
attn_metadata: PerLayerAttnMetadata | None = None
# If force_attention is True, we always capture attention. Otherwise,
# it only happens for cudagraph_runtime_mode=FULL.
@@ -405,41 +386,30 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
self.query_start_loc.copy_to_gpu()
for kv_cache_group_id, kv_cache_group_spec in enumerate(
self.kv_cache_config.kv_cache_groups
):
for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups):
common_attn_metadata = CommonAttentionMetadata(
query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
seq_lens=self.seq_lens.gpu[:num_reqs],
seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[ # noqa: E501
:num_reqs
],
num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs], # noqa: E501
num_reqs=num_reqs,
num_actual_tokens=num_tokens,
max_query_len=max_query_len,
max_seq_len=self.max_model_len,
block_table_tensor=self.input_batch.block_table[
kv_cache_group_id
].get_device_tensor(num_reqs),
slot_mapping=self.input_batch.block_table[
kv_cache_group_id
].slot_mapping.gpu[:num_tokens],
block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs),
slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens],
causal=True,
)
for attn_group in self.attn_groups[kv_cache_group_id]:
assert type(attn_metadata) is dict
attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture( # noqa: E501
attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(
common_attn_metadata
)
) # noqa: E501
for layer_name in attn_group.layer_names:
attn_metadata[layer_name] = attn_metadata_i
with self.maybe_dummy_run_with_lora(
self.lora_config, num_scheduled_tokens, remove_lora
):
with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora):
model_kwargs = self._init_model_kwargs(num_tokens)
if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
input_ids = None
@@ -465,17 +435,13 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
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,
)
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
)
intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens, None, False)
# filter out the valid batch descriptor
_cg_mode, batch_descriptor = (
@@ -491,10 +457,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
if cudagraph_runtime_mode is not None:
# we allow forcing NONE when the dispatcher disagrees to support
# warm ups for cudagraph capture
assert (
cudagraph_runtime_mode == CUDAGraphMode.NONE
or cudagraph_runtime_mode == _cg_mode
), (
assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, (
f"Cudagraph runtime mode mismatch at dummy_run. "
f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
)
@@ -550,10 +513,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
# Profile with multimodal encoder & encoder cache.
if self.supports_mm_inputs:
if self.model_config.multimodal_config.skip_mm_profiling:
logger.info(
"Skipping memory profiling for multimodal encoder and "
"encoder cache."
)
logger.info("Skipping memory profiling for multimodal encoder and encoder cache.")
else:
mm_budget = self.mm_budget
assert mm_budget is not None
@@ -563,14 +523,11 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
# modality with the max possible input tokens even when
# it supports multiple.
dummy_modality = mm_budget.get_modality_with_max_tokens()
max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[
dummy_modality
]
max_mm_items_per_batch = mm_budget.max_items_per_batch_by_modality[dummy_modality]
logger.info(
"Encoder cache will be initialized with a budget of "
"%s tokens, and profiled with %s %s items of the "
"maximum feature size.",
"Encoder cache will be initialized with a budget "
"of %s tokens, and profiled with %s %s items of the maximum feature size.",
encoder_budget,
max_mm_items_per_batch,
dummy_modality,
@@ -583,9 +540,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
)
# Run multimodal encoder.
dummy_encoder_outputs = self.model.get_multimodal_embeddings(
**batched_dummy_mm_inputs
)
dummy_encoder_outputs = self.model.get_multimodal_embeddings(**batched_dummy_mm_inputs)
sanity_check_mm_encoder_outputs(
dummy_encoder_outputs,
@@ -601,9 +556,7 @@ class GPUDiffusionModelRunner(OmniGPUModelRunner):
if encoder_output_shape[0] < encoder_budget:
expanded_outputs = []
for output in dummy_encoder_outputs:
expanded = output.new_zeros(
(encoder_budget, encoder_output_shape[-1])
)
expanded = output.new_zeros((encoder_budget, encoder_output_shape[-1]))
num_tokens = output.shape[0]
expanded[:num_tokens].copy_(output)
expanded_outputs.append(expanded)
vllm_omni/worker/gpu_diffusion_worker.py
+3 -8
View File
@@ -2,13 +2,13 @@ import gc
import os
import torch
from vllm.model_executor import set_random_seed
from vllm.platforms import current_platform
from vllm.utils import GiB_bytes, MemorySnapshot
from vllm.v1.utils import report_usage_stats
from vllm.v1.worker.gpu_worker import Worker as GPUWorker
from vllm.v1.worker.gpu_worker import init_worker_distributed_environment
from vllm_omni.worker.gpu_diffusion_model_runner import GPUDiffusionModelRunner
@@ -43,10 +43,7 @@ class GPUDiffusionWorker(GPUWorker):
# take current memory snapshot
self.init_snapshot = MemorySnapshot()
self.requested_memory = (
self.init_snapshot.total_memory
* self.cache_config.gpu_memory_utilization
)
self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization
if self.init_snapshot.free_memory < self.requested_memory:
def GiB(b):
@@ -65,9 +62,7 @@ class GPUDiffusionWorker(GPUWorker):
raise RuntimeError(f"Not support device type: {self.device_config.device}")
# Construct the model runner
self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(
self.vllm_config, self.device
)
self.model_runner: GPUDiffusionModelRunner = GPUDiffusionModelRunner(self.vllm_config, self.device)
if self.rank == 0:
# If usage stat is enabled, collect relevant info.
vllm_omni/worker/gpu_model_runner.py
+27 -72
View File
@@ -1,8 +1,7 @@
from typing import TYPE_CHECKING, List, Optional, Union, cast
from typing import TYPE_CHECKING, Optional, Union, cast
import numpy as np
import torch
from vllm.config import CUDAGraphMode
from vllm.distributed.parallel_state import get_pp_group
from vllm.forward_context import BatchDescriptor, set_forward_context
@@ -10,10 +9,7 @@ from vllm.logger import init_logger
from vllm.model_executor.models.interfaces_base import VllmModelForPooling
from vllm.sampling_params import SamplingType
from vllm.utils import LazyLoader, cdiv
from vllm.v1.attention.backends.utils import (
CommonAttentionMetadata,
split_attn_metadata,
)
from vllm.v1.attention.backends.utils import CommonAttentionMetadata, split_attn_metadata
from vllm.v1.spec_decode.eagle import EagleProposer
from vllm.v1.worker.gpu_input_batch import CachedRequestState
from vllm.v1.worker.gpu_model_runner import GPUModelRunner, PerLayerAttnMetadata
@@ -81,10 +77,7 @@ class OmniGPUModelRunner(GPUModelRunner):
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
):
if sampling_params and sampling_params.sampling_type == SamplingType.RANDOM_SEED:
generator = torch.Generator(device=self.device)
generator.manual_seed(sampling_params.seed)
else:
@@ -145,9 +138,7 @@ class OmniGPUModelRunner(GPUModelRunner):
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")
)
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)
@@ -188,9 +179,7 @@ class OmniGPUModelRunner(GPUModelRunner):
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
)
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])
@@ -228,23 +217,17 @@ class OmniGPUModelRunner(GPUModelRunner):
# 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.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, ()
)
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
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
@@ -261,20 +244,15 @@ class OmniGPUModelRunner(GPUModelRunner):
self.input_batch.refresh_metadata()
@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
):
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):
elif isinstance(hidden_states, list):
text_hidden_states = hidden_states[0]
multimodal_outputs = {}
else:
@@ -427,37 +405,25 @@ class OmniGPUModelRunner(GPUModelRunner):
self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens
self.query_start_loc.copy_to_gpu()
for kv_cache_group_id, kv_cache_group_spec in enumerate(
self.kv_cache_config.kv_cache_groups
):
for kv_cache_group_id, kv_cache_group_spec in enumerate(self.kv_cache_config.kv_cache_groups):
common_attn_metadata = CommonAttentionMetadata(
query_start_loc=self.query_start_loc.gpu[: num_reqs + 1],
query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1],
seq_lens=self.seq_lens.gpu[:num_reqs],
seq_lens_cpu=self.seq_lens.cpu[:num_reqs],
num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[ # noqa: E501
:num_reqs
],
num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs], # noqa: E501
num_reqs=num_reqs,
num_actual_tokens=num_tokens,
max_query_len=max_query_len,
max_seq_len=self.max_model_len,
block_table_tensor=self.input_batch.block_table[
kv_cache_group_id
].get_device_tensor(num_reqs),
slot_mapping=self.input_batch.block_table[
kv_cache_group_id
].slot_mapping.gpu[:num_tokens],
block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs),
slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens],
causal=True,
)
for attn_group in self.attn_groups[kv_cache_group_id]:
if ubatch_slices is not None:
common_attn_metadata_list = split_attn_metadata(
ubatch_slices, common_attn_metadata
)
for ubid, common_attn_metadata in enumerate(
common_attn_metadata_list
):
common_attn_metadata_list = split_attn_metadata(ubatch_slices, common_attn_metadata)
for ubid, common_attn_metadata in enumerate(common_attn_metadata_list):
assert common_attn_metadata.max_query_len == 1
attn_metadata_i = attn_group.get_metadata_builder(
ubatch_id=ubid
@@ -467,15 +433,13 @@ class OmniGPUModelRunner(GPUModelRunner):
attn_metadata[ubid][layer_name] = attn_metadata_i
else:
assert type(attn_metadata) is dict
attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture( # noqa: E501
attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture(
common_attn_metadata
)
) # noqa: E501
for layer_name in attn_group.layer_names:
attn_metadata[layer_name] = attn_metadata_i
with self.maybe_dummy_run_with_lora(
self.lora_config, num_scheduled_tokens, remove_lora
):
with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora):
model_kwargs = self._init_model_kwargs(num_tokens)
if self.supports_mm_inputs and not self.model_config.is_encoder_decoder:
input_ids = None
@@ -501,17 +465,13 @@ class OmniGPUModelRunner(GPUModelRunner):
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,
)
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
)
intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens, None, False)
# filter out the valid batch descriptor
_cg_mode, batch_descriptor = (
@@ -527,10 +487,7 @@ class OmniGPUModelRunner(GPUModelRunner):
if cudagraph_runtime_mode is not None:
# we allow forcing NONE when the dispatcher disagrees to support
# warm ups for cudagraph capture
assert (
cudagraph_runtime_mode == CUDAGraphMode.NONE
or cudagraph_runtime_mode == _cg_mode
), (
assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, (
f"Cudagraph runtime mode mismatch at dummy_run. "
f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}."
)
@@ -585,7 +542,5 @@ class OmniGPUModelRunner(GPUModelRunner):
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
)
hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states)
return hidden_states, hidden_states[logit_indices]