diff --git a/.github/workflows/ci-all-files.yml b/.github/workflows/ci-all-files.yml index 01420ef54..a5e4179a0 100644 --- a/.github/workflows/ci-all-files.yml +++ b/.github/workflows/ci-all-files.yml @@ -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]" diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2eaab59fa..413c19092 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -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' diff --git a/.github/workflows/pre-commit.yml b/.github/workflows/pre-commit.yml index 5decd423c..23292bdbd 100644 --- a/.github/workflows/pre-commit.yml +++ b/.github/workflows/pre-commit.yml @@ -22,5 +22,3 @@ jobs: python-version: "3.12" cache: 'pip' - uses: pre-commit/action@v3.0.1 - with: - extra_args: --hook-stage manual diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 58f7c1373..861d131bb 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -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 diff --git a/README.md b/README.md index e8a345c5e..45f23d42c 100644 --- a/README.md +++ b/README.md @@ -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" ``` diff --git a/docs/contributing/design_documents/vllm_omni_design.md b/docs/contributing/design_documents/vllm_omni_design.md index 6275c8950..27b62f691 100644 --- a/docs/contributing/design_documents/vllm_omni_design.md +++ b/docs/contributing/design_documents/vllm_omni_design.md @@ -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, diff --git a/docs/mkdocs/hooks/generate_examples.py b/docs/mkdocs/hooks/generate_examples.py index 82b1ba06f..a1f75983f 100644 --- a/docs/mkdocs/hooks/generate_examples.py +++ b/docs/mkdocs/hooks/generate_examples.py @@ -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 diff --git a/docs/user_guide/examples/offline_inference/qwen_2_5_omni.md b/docs/user_guide/examples/offline_inference/qwen_2_5_omni.md index 38e2f327a..0f9aa971f 100644 --- a/docs/user_guide/examples/offline_inference/qwen_2_5_omni.md +++ b/docs/user_guide/examples/offline_inference/qwen_2_5_omni.md @@ -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" ``` diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index f205afcbb..e3e496f46 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -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 (with --prompt_type text)" - ) + raise ValueError("No prompts provided. Use --prompts ... or --txt-prompts (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}") diff --git a/examples/offline_inference/qwen2_5_omni/extract_prompts.py b/examples/offline_inference/qwen2_5_omni/extract_prompts.py index cf39f46a2..3d940b995 100644 --- a/examples/offline_inference/qwen2_5_omni/extract_prompts.py +++ b/examples/offline_inference/qwen2_5_omni/extract_prompts.py @@ -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 diff --git a/examples/offline_inference/qwen2_5_omni/processing_omni.py b/examples/offline_inference/qwen2_5_omni/processing_omni.py index f379b6c38..a22220dd3 100644 --- a/examples/offline_inference/qwen2_5_omni/processing_omni.py +++ b/examples/offline_inference/qwen2_5_omni/processing_omni.py @@ -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) diff --git a/examples/offline_inference/qwen2_5_omni/utils.py b/examples/offline_inference/qwen2_5_omni/utils.py index c25b0734f..11df5c55a 100644 --- a/examples/offline_inference/qwen2_5_omni/utils.py +++ b/examples/offline_inference/qwen2_5_omni/utils.py @@ -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": diff --git a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py index 3bc864852..13d37476d 100644 --- a/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py +++ b/examples/online_serving/openai_chat_completion_client_for_multimodal_generation.py @@ -16,7 +16,6 @@ SEED = 42 def get_system_prompt(): - return { "role": "system", "content": [ diff --git a/pyproject.toml b/pyproject.toml index 5f0e38f8a..95ab9111d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -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" +] diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 00e56813d..67bcc5fad 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -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 diff --git a/vllm_omni/core/sched/diffusion_scheduler.py b/vllm_omni/core/sched/diffusion_scheduler.py index 8aa7431a1..0474f83ab 100644 --- a/vllm_omni/core/sched/diffusion_scheduler.py +++ b/vllm_omni/core/sched/diffusion_scheduler.py @@ -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 diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index 45eb6bee0..dbfc20019 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -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 diff --git a/vllm_omni/core/sched/scheduler.py b/vllm_omni/core/sched/scheduler.py index fbd02c4b7..3a30543c7 100644 --- a/vllm_omni/core/sched/scheduler.py +++ b/vllm_omni/core/sched/scheduler.py @@ -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 diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index dd3a9a710..2de831388 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -5,7 +5,6 @@ Engine components for vLLM-omni. from typing import Any, Optional import msgspec - from vllm.v1.engine import EngineCoreRequest diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 62f7b1343..64df79c7b 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -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 diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index f870a0f16..fb664b7a2 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -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 diff --git a/vllm_omni/engine/processor.py b/vllm_omni/engine/processor.py index d74022535..c25d7da0b 100644 --- a/vllm_omni/engine/processor.py +++ b/vllm_omni/engine/processor.py @@ -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, diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py index cc1114afc..cfde9e14a 100644 --- a/vllm_omni/entrypoints/async_omni_llm.py +++ b/vllm_omni/entrypoints/async_omni_llm.py @@ -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, diff --git a/vllm_omni/entrypoints/cli/main.py b/vllm_omni/entrypoints/cli/main.py index 3fb8bf188..148839d75 100644 --- a/vllm_omni/entrypoints/cli/main.py +++ b/vllm_omni/entrypoints/cli/main.py @@ -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, ] diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 5fda170fa..25746cb5e 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -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, diff --git a/vllm_omni/entrypoints/log_utils.py b/vllm_omni/entrypoints/log_utils.py index 8f616db0f..00172bcaf 100644 --- a/vllm_omni/entrypoints/log_utils.py +++ b/vllm_omni/entrypoints/log_utils.py @@ -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), diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 2eb493e80..655d0307f 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -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.log" - ) + suggestions.append(f"Inspect per-stage log files for details: {self._log_file}.stage.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) diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index bbb33c41d..994dfa423 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -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, diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 4c3931466..da403d804 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -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()) diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index b4f2bcb59..6e7ade7c7 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -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) diff --git a/vllm_omni/entrypoints/stage_utils.py b/vllm_omni/entrypoints/stage_utils.py index dc38e617a..36ad2b0ab 100644 --- a/vllm_omni/entrypoints/stage_utils.py +++ b/vllm_omni/entrypoints/stage_utils.py @@ -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) diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index a0de525bf..7662304c0 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -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__) diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 6feff5940..7c518bcc8 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -7,7 +7,6 @@ except ImportError: from typing_extensions import NotRequired import torch - from vllm.inputs.data import EmbedsPrompt, TokenInputs, TokensPrompt diff --git a/vllm_omni/inputs/parse.py b/vllm_omni/inputs/parse.py index 9bf505779..ace3b7a44 100644 --- a/vllm_omni/inputs/parse.py +++ b/vllm_omni/inputs/parse.py @@ -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") diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 31c491ae3..28046b802 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -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") diff --git a/vllm_omni/model_executor/layers/mrope.py b/vllm_omni/model_executor/layers/mrope.py index e986c2b8d..789b3b3cf 100644 --- a/vllm_omni/model_executor/layers/mrope.py +++ b/vllm_omni/model_executor/layers/mrope.py @@ -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) diff --git a/vllm_omni/model_executor/model_loader/weight_utils.py b/vllm_omni/model_executor/model_loader/weight_utils.py index 949c858dc..f20ee6b05 100644 --- a/vllm_omni/model_executor/model_loader/weight_utils.py +++ b/vllm_omni/model_executor/model_loader/weight_utils.py @@ -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 diff --git a/vllm_omni/model_executor/models/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni.py index f6e2ab7fc..ae5ff25ca 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni.py @@ -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) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni_talker.py index 8944b693b..4d906f79a 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni_talker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni_talker.py @@ -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 [] diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py index f1b20cfeb..41036a567 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni_thinker.py @@ -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 ), diff --git a/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py b/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py index d2271ba1f..f1109a9a7 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni_token2wav.py @@ -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 diff --git a/vllm_omni/model_executor/models/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_old.py index 0ffe22b22..c0145860b 100644 --- a/vllm_omni/model_executor/models/qwen2_old.py +++ b/vllm_omni/model_executor/models/qwen2_old.py @@ -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) diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index b01e17fa0..25b16c6a5 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -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": ( diff --git a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py index 16717cdfc..c1a04d9b0 100644 --- a/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py +++ b/vllm_omni/model_executor/stage_input_processors/qwen2_5_omni.py @@ -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, ) diff --git a/vllm_omni/outputs.py b/vllm_omni/outputs.py index cf97114eb..454e7b423 100644 --- a/vllm_omni/outputs.py +++ b/vllm_omni/outputs.py @@ -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 diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index b8de3010b..ea0828cff 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -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 diff --git a/vllm_omni/request.py b/vllm_omni/request.py index cb5e012f5..bc71bf968 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -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, diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 8f757bb6a..7b793214b 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -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, ) diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 838068df7..98826a93e 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -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. diff --git a/vllm_omni/worker/gpu_diffusion_model_runner.py b/vllm_omni/worker/gpu_diffusion_model_runner.py index c8b294c5d..d381db8d1 100644 --- a/vllm_omni/worker/gpu_diffusion_model_runner.py +++ b/vllm_omni/worker/gpu_diffusion_model_runner.py @@ -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) diff --git a/vllm_omni/worker/gpu_diffusion_worker.py b/vllm_omni/worker/gpu_diffusion_worker.py index 46af50190..77f1f03fb 100644 --- a/vllm_omni/worker/gpu_diffusion_worker.py +++ b/vllm_omni/worker/gpu_diffusion_worker.py @@ -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. diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 67dfdf537..0a90dbdb8 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -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]