diff --git a/.buildkite/scripts/simple_test.sh b/.buildkite/scripts/simple_test.sh index 157ad9ca0..bb2e1eca0 100755 --- a/.buildkite/scripts/simple_test.sh +++ b/.buildkite/scripts/simple_test.sh @@ -47,7 +47,7 @@ fi VENV_PYTHON="${VENV_DIR}/bin/python" [[ -x "${VENV_PYTHON}" ]] || { echo "Python not found in ${VENV_DIR}"; exit 1; } -"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.11.0 +"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.12.0 "${UV_BIN}" pip install --python "${VENV_PYTHON}" -e ".[dev]" "${VENV_PYTHON}" -m pytest -v -s tests/entrypoints/ "${VENV_PYTHON}" -m pytest -v -s tests/diffusion/cache/ diff --git a/.gitignore b/.gitignore index 92ed5ad9a..f7b771bc7 100644 --- a/.gitignore +++ b/.gitignore @@ -221,6 +221,7 @@ datasets/ outputs/ results/ generated/ +output_*/ # Configuration overrides configs/local.yaml diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 0d513c054..5e1d00a5f 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -1,5 +1,5 @@ ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.11.0 +ARG VLLM_BASE_TAG=v0.12.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index 52d365d9c..e7df09ff4 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -18,9 +18,9 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** #### Installation of vLLM -vLLM-Omni is built based on vLLM v0.11.0. Please install it with command below. +vLLM-Omni is built based on vLLM v0.12.0. Please install it with command below. ```bash -uv pip install vllm==0.11.0 --torch-backend=auto +uv pip install vllm==0.12.0 --torch-backend=auto ``` #### Installation of vLLM-Omni @@ -34,10 +34,10 @@ uv pip install vllm-omni # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.11.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.12.0 release version of the library ```bash -uv pip install vllm==0.11.0 --torch-backend=auto +uv pip install vllm==0.12.0 --torch-backend=auto ``` #### Installation of vLLM-Omni @@ -54,11 +54,11 @@ If you want to check, modify or debug with source code of vLLM, install the libr ```bash git clone https://github.com/vllm-project/vllm.git cd vllm -git checkout v0.11.0 +git checkout v0.12.0 ``` Set up environment variables to get pre-built wheels. If there are internet problems, just download the whl file manually. And set `VLLM_PRECOMPILED_WHEEL_LOCATION` as your local absolute path of whl file. ```bash -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.11.0/vllm-0.11.0-cp38-abi3-manylinux1_x86_64.whl +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.12.0/vllm-0.12.0-cp38-abi3-manylinux_2_31_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash @@ -89,7 +89,7 @@ docker run --runtime nvidia --gpus 2 \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ --ipc=host \ - vllm/vllm-omni:v0.11.0rc1 \ + vllm/vllm-omni:v0.12.0rc1 \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct --port 8091 ``` diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index 3b970267a..87d7eb27a 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -19,7 +19,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images. #### Get the correct docker environment -The v0.11.0 releases is on Oct 3, 2025, so we will pull a AMD nightly docker image close to the release period. `rocm/vllm-dev:nightly_main_20251005` is picked as it has a correct `transformers` version that supports qwen3_omni. +The v0.12.0 releases is on Dec 3, 2025, so we will pull a AMD nightly docker image close to the release period. `rocm/vllm-dev:nightly_main_20251205` is picked as it has a correct `transformers` version that supports qwen3_omni. `docker pull rocm/vllm-dev:nightly_main_20251005` @@ -51,18 +51,18 @@ sudo apt install ffmpeg -y ``` #### Installation of vLLM -Install the v0.11.0 release of vLLM as vLLM-Omni is currently depend on v0.11.0. +Install the v0.12.0 release of vLLM as vLLM-Omni is currently depend on v0.12.0. ```bash # uninstall existing vllm which is based on the main branch, not on the releases branch python3 -m pip uninstall -y vllm -git clone https://github.com/vllm-project/vllm.git vllm0110 -cd vllm0110/ +git clone https://github.com/vllm-project/vllm.git vllm0120 +cd vllm0120/ -# checkout the version tag v0.11.0 -git checkout -b v0.11.0 v0.11.0 +# checkout the version tag v0.12.0 +git checkout -b v0.12.0 v0.12.0 -# you should see 0.11.0 +# you should see 0.12.0 python -c "import setuptools_scm; print(setuptools_scm.get_version())" # PYTORCH_ROCM_ARCH= python3 setup.py develop @@ -92,7 +92,7 @@ pip install vllm-omni # clone and checkout to the release version. git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni -git checkout v0.11.0rc1 +git checkout v0.12.0rc1 pip install -e . ``` diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index 0b18f77d4..90ee150f0 100644 --- a/docs/getting_started/installation/npu/npu.inc.md +++ b/docs/getting_started/installation/npu/npu.inc.md @@ -13,10 +13,10 @@ export DEVICE0=/dev/davinci0 export DEVICE1=/dev/davinci1 # Update the vllm-ascend image # Atlas A2: -# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2 +# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1 # Atlas A3: -# export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2-a3 -export IMAGE=quay.io/ascend/vllm-ascend:v0.11.0rc2 +# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1-a3 +export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1 docker run --rm \ --name vllm-omni-npu \ --shm-size=1g \ diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index a2a47bf3e..bfd5324c5 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -17,7 +17,7 @@ from vllm.assets.image import ImageAsset from vllm.assets.video import VideoAsset, video_to_ndarrays from vllm.multimodal.image import convert_image_mode from vllm.sampling_params import SamplingParams -from vllm.utils import FlexibleArgumentParser +from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm_omni.entrypoints.omni import Omni @@ -119,7 +119,7 @@ def get_use_audio_in_video_query( question = "Describe the content of the video, then convert what the baby say into text." prompt = ( f"<|im_start|>system\n{default_system}<|im_end|>\n" - "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|>" + "<|im_start|>user\n<|vision_bos|><|VIDEO|><|vision_eos|><|audio_bos|><|AUDIO|><|audio_eos|>" f"{question}<|im_end|>\n" f"<|im_start|>assistant\n" ) @@ -319,7 +319,6 @@ def main(args): query_result = query_func(audio_path=audio_path, sampling_rate=sampling_rate) else: query_result = query_func() - omni_llm = Omni( model=model_name, log_stats=args.enable_stats, diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index eb580bcf9..753f7cc36 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -17,7 +17,7 @@ from vllm.assets.audio import AudioAsset from vllm.assets.image import ImageAsset from vllm.assets.video import VideoAsset, video_to_ndarrays from vllm.multimodal.image import convert_image_mode -from vllm.utils import FlexibleArgumentParser +from vllm.utils.argparse_utils import FlexibleArgumentParser from vllm_omni.entrypoints.omni import Omni @@ -281,7 +281,7 @@ def parse_args(): "--query-type", "-q", type=str, - default="mixed_modalities", + default="use_video", choices=query_map.keys(), help="Query type.", ) diff --git a/examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py index e5d430d48..e5ef42c8c 100644 --- a/examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py +++ b/examples/online_serving/qwen2_5_omni/openai_chat_completion_client_for_multimodal_generation.py @@ -4,7 +4,7 @@ import os import requests from openai import OpenAI from vllm.assets.audio import AudioAsset -from vllm.utils import FlexibleArgumentParser +from vllm.utils.argparse_utils import FlexibleArgumentParser # Modify OpenAI's API key and API base to use vLLM's API server. openai_api_key = "EMPTY" diff --git a/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py b/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py index 4a31ddfa2..7ebf8c7f4 100644 --- a/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py +++ b/examples/online_serving/qwen3_omni/openai_chat_completion_client_for_multimodal_generation.py @@ -5,7 +5,7 @@ from typing import NamedTuple import requests from openai import OpenAI from vllm.assets.audio import AudioAsset -from vllm.utils import FlexibleArgumentParser +from vllm.utils.argparse_utils import FlexibleArgumentParser # Modify OpenAI's API key and API base to use vLLM's API server. openai_api_key = "EMPTY" diff --git a/pyproject.toml b/pyproject.toml index 2aebde92e..34eca4f32 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta" [project] name = "vllm-omni" -version = "0.11.0rc1" +version = "0.12.0rc1" description = "A framework for efficient model inference with omni-modality models" readme = "README.md" requires-python = ">=3.10,<3.14" @@ -37,7 +37,7 @@ dependencies = [ "gradio==5.50", "soundfile>=0.13.1", "cache-dit==1.1.8", - # "vllm==0.11.0", # TODO: fix the entrypoints overwrite problem + # "vllm==0.12.0", # TODO: fix the entrypoints overwrite problem ] [project.optional-dependencies] diff --git a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml index e9f87be38..bcaea6703 100644 --- a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml @@ -44,7 +44,7 @@ stage_args: model_arch: Qwen3OmniMoeForConditionalGeneration worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler - gpu_memory_utilization: 0.5 + gpu_memory_utilization: 0.32 enforce_eager: true trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index b43fa8361..945b6eaef 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -27,7 +27,7 @@ test_params = [(model, stage_config) for model in models for stage_config in sta def test_video_to_audio(omni_runner: type[OmniRunner], test_config) -> None: """Test processing video, generating audio output.""" model, stage_config_path = test_config - with omni_runner(model, seed=42, stage_configs_path=stage_config_path) as runner: + with omni_runner(model, seed=42, stage_configs_path=stage_config_path, init_sleep_seconds=90) as runner: # Prepare inputs question = "Describe the video briefly." video = VideoAsset(name="baby_reading", num_frames=4).np_ndarrays diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index 147bcb07c..a0f287294 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -109,7 +109,7 @@ def omni_server(request): Multi-stage initialization can take 10-20+ minutes. """ model, stage_config_path = request.param - with OmniServer(model, ["--stage-configs-path", stage_config_path]) as server: + with OmniServer(model, ["--stage-configs-path", stage_config_path, "--init-sleep-seconds", "90"]) as server: yield server diff --git a/tests/entrypoints/test_omni_llm.py b/tests/entrypoints/test_omni_llm.py index e2fd12427..a3cb236de 100644 --- a/tests/entrypoints/test_omni_llm.py +++ b/tests/entrypoints/test_omni_llm.py @@ -393,12 +393,12 @@ def mock_get_config(monkeypatch): ) # Also mock in processor module if it's imported monkeypatch.setattr( - "vllm_omni.engine.processor.length_from_prompt_token_ids_or_embeds", + "vllm_omni.engine.input_processor.length_from_prompt_token_ids_or_embeds", _mock_length_from_prompt_token_ids_or_embeds, raising=False, ) # If processor module is already imported, patch it directly - processor_module_path = "vllm_omni.engine.processor" + processor_module_path = "vllm_omni.engine.input_processor" if processor_module_path in sys.modules: processor_module = sys.modules[processor_module_path] setattr( diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index 8fc929c38..8222c3076 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -1,17 +1,17 @@ -import json import warnings from importlib.util import find_spec -from typing import Any, Literal +from typing import Any import torch import vllm.envs as envs -from omegaconf import DictConfig, OmegaConf from pydantic import ConfigDict from pydantic.dataclasses import dataclass +from vllm.attention.backends.registry import AttentionBackendEnum from vllm.config import ModelConfig, config from vllm.config.model import ( _RUNNER_CONVERTS, _RUNNER_TASKS, + ConvertOption, ConvertType, RunnerOption, TaskOption, @@ -27,9 +27,14 @@ from vllm.transformers_utils.config import ( get_hf_image_processor_config, get_hf_text_config, get_pooling_config, - is_interleaved, ) -from vllm.transformers_utils.utils import maybe_model_redirect +from vllm.transformers_utils.gguf_utils import ( + maybe_patch_hf_config_from_gguf, +) +from vllm.transformers_utils.utils import ( + is_gguf, + maybe_model_redirect, +) import vllm_omni.model_executor.models as me_models @@ -87,29 +92,19 @@ class OmniModelConfig(ModelConfig): def __post_init__( self, # Multimodal config init vars - limit_mm_per_prompt: dict[str, int] | None, + limit_mm_per_prompt: dict[str, int | dict[str, int]] | None, + enable_mm_embeds: bool | None, media_io_kwargs: dict[str, dict[str, Any]] | None, mm_processor_kwargs: dict[str, Any] | None, mm_processor_cache_gb: float | None, mm_processor_cache_type: MMCacheType | None, mm_shm_cache_max_object_size_mb: int | None, mm_encoder_tp_mode: MMEncoderTPMode | None, + mm_encoder_attn_backend: AttentionBackendEnum | str | None, interleave_mm_strings: bool | None, skip_mm_profiling: bool | None, video_pruning_rate: float | None, ) -> None: - # Set the default seed to 0 in V1. - if envs.VLLM_USE_V1 and self.seed is None: - self.seed = 0 - if not envs.VLLM_ENABLE_V1_MULTIPROCESSING: - logger.warning( - "The global random seed is set to %d. Since " - "VLLM_ENABLE_V1_MULTIPROCESSING is set to False, this may " - "affect the random state of the Python process that " - "launched vLLM.", - self.seed, - ) - # Keep set served_model_name before maybe_model_redirect(self.model) self.served_model_name = get_served_model_name(self.model, self.served_model_name) self.model = maybe_model_redirect(self.model) @@ -124,53 +119,36 @@ class OmniModelConfig(ModelConfig): self.hf_config_path = maybe_model_redirect(self.hf_config_path) if callable(self.hf_overrides): - hf_overrides_kw: dict[str, Any] = {} + hf_overrides_kw = {} hf_overrides_fn = self.hf_overrides dict_overrides: dict[str, Any] = {} else: - # Separate dict overrides from flat ones; convert OmegaConf to dict + # Separate dict overrides from flat ones + # We'll determine how to apply dict overrides after loading the config hf_overrides_kw = {} dict_overrides = {} for key, value in self.hf_overrides.items(): - if isinstance(value, DictConfig): - value = OmegaConf.to_container(value, resolve=True, structured_config_mode=None) if isinstance(value, dict): dict_overrides[key] = value else: hf_overrides_kw[key] = value hf_overrides_fn = None - if self.rope_scaling: - hf_override: dict[str, Any] = {"rope_scaling": self.rope_scaling} - hf_overrides_kw.update(hf_override) - hf_overrides_str = json.dumps(hf_overrides_kw) - msg = ( - "`--rope-scaling` will be removed in a future release. " - f"'Please instead use `--hf-overrides '{hf_overrides_str}'`" - ) - warnings.warn(DeprecationWarning(msg), stacklevel=2) - if self.rope_theta is not None: - hf_override = {"rope_theta": self.rope_theta} - hf_overrides_kw.update(hf_override) - hf_overrides_str = json.dumps(hf_overrides_kw) - msg = ( - "`--rope-theta` will be removed in a future release. " - f"'Please instead use `--hf-overrides '{hf_overrides_str}'`" - ) - warnings.warn(DeprecationWarning(msg), stacklevel=2) - self.maybe_pull_model_tokenizer_for_runai(self.model, self.tokenizer) if (backend := envs.VLLM_ATTENTION_BACKEND) and backend == "FLASHINFER" and find_spec("flashinfer") is None: raise ValueError( "VLLM_ATTENTION_BACKEND is set to FLASHINFER, but flashinfer " "module was not found. See " - "https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile " + "https://github.com/vllm-project/vllm/blob/main/docker/Dockerfile " # noqa: E501 "for instructions on how to install it." ) if self.override_attention_dtype is not None and not current_platform.is_rocm(): - warnings.warn("override-attention-dtype is set but not using ROCm platform", stacklevel=2) + warnings.warn( + "override-attention-dtype is set but not using ROCm platform", + stacklevel=2, + ) if self.enable_sleep_mode and not current_platform.is_sleep_mode_available(): raise ValueError("Sleep mode is not supported on current platform.") @@ -184,11 +162,14 @@ class OmniModelConfig(ModelConfig): hf_overrides_kw=hf_overrides_kw, hf_overrides_fn=hf_overrides_fn, ) - - if dict_overrides: - self._apply_dict_overrides(hf_config, dict_overrides) + hf_config = maybe_patch_hf_config_from_gguf( + self.model, + hf_config, + ) self.hf_config = hf_config + if dict_overrides: + self._apply_dict_overrides(hf_config, dict_overrides) self.hf_text_config = self.draw_hf_text_config() self.attention_chunk_size = getattr(self.hf_text_config, "attention_chunk_size", None) self.encoder_config = self._get_encoder_config() @@ -216,7 +197,7 @@ class OmniModelConfig(ModelConfig): if self.task is not None: runner: RunnerOption = "auto" - convert: ConvertType | Literal["auto"] = "auto" # type: ignore[name-defined] + convert: ConvertOption = "auto" msg_prefix = ( "The 'task' option has been deprecated and will be removed in v0.13.0 or v1.0, whichever comes first." ) @@ -228,7 +209,7 @@ class OmniModelConfig(ModelConfig): if is_generative_model and is_pooling_model: if is_generative_task: runner = "generate" - convert = "auto" # type: ignore[assignment] + convert = "auto" msg_hint = ( "Please replace this option with `--runner " "generate` to continue using this model " @@ -236,7 +217,7 @@ class OmniModelConfig(ModelConfig): ) elif is_pooling_task: runner = "pooling" - convert = "auto" # type: ignore[assignment] + convert = "auto" msg_hint = ( "Please replace this option with `--runner " "pooling` to continue using this model " @@ -247,7 +228,7 @@ class OmniModelConfig(ModelConfig): elif is_generative_model or is_pooling_model: if is_generative_task: runner = "generate" - convert = "auto" # type: ignore[assignment] + convert = "auto" msg_hint = "Please remove this option" elif is_pooling_task: runner = "pooling" @@ -260,12 +241,29 @@ class OmniModelConfig(ModelConfig): else: # task == "auto" pass else: - raise AssertionError( - f"The model should be a generative or pooling model when task is set to {self.task!r}." - ) + # Neither generative nor pooling model - try to convert if possible + if is_pooling_task: + runner = "pooling" + convert = _task_to_convert(self.task) + msg_hint = ( + "Please replace this option with `--runner pooling " + f"--convert {convert}` to continue using this model " + "as a pooling model." + ) + else: + debug_info = { + "architectures": architectures, + "is_generative_model": is_generative_model, + "is_pooling_model": is_pooling_model, + } + raise AssertionError( + "The model should be a generative or " + "pooling model when task is set to " + f"{self.task!r}. Found: {debug_info}" + ) self.runner = runner - self.convert = convert # type: ignore[assignment] + self.convert = convert msg = f"{msg_prefix} {msg_hint}" warnings.warn(msg, DeprecationWarning, stacklevel=2) @@ -276,6 +274,7 @@ class OmniModelConfig(ModelConfig): if self.runner_type == "generate" and not is_generative_model: generate_converts = _RUNNER_CONVERTS["generate"] if self.convert_type not in generate_converts: + # Currently we don't have any converters for generative models raise ValueError("This model does not support `--runner generate`.") if self.runner_type == "pooling" and not is_pooling_model: pooling_converts = _RUNNER_CONVERTS["pooling"] @@ -296,23 +295,12 @@ class OmniModelConfig(ModelConfig): # Init pooler config if needed if self.runner_type == "pooling": - if self.override_pooler_config is not None: - logger.warning_once( - "`override_pooler_config` is deprecated and will be " - "removed in v0.12.0 or v1.0.0, whichever is sooner. " - "Please use `pooler_config` instead." - ) - - if isinstance(self.override_pooler_config, dict): - self.pooler_config = PoolerConfig(**self.override_pooler_config) - else: - self.pooler_config = self.override_pooler_config - if self.pooler_config is None: self.pooler_config = PoolerConfig() base_config = get_pooling_config(self.model, self.revision) if base_config is not None: + # Only set values that are not overridden by the user for k, v in base_config.items(): if getattr(self.pooler_config, k) is None: setattr(self.pooler_config, k, v) @@ -329,23 +317,6 @@ class OmniModelConfig(ModelConfig): revision=self.revision, ) - # Interleaved attention is not supported by some backends in V0 - if ( - not self.disable_sliding_window - and is_interleaved(self.hf_text_config) - and not envs.VLLM_USE_V1 - and (backend := envs.VLLM_ATTENTION_BACKEND) in ("XFORMERS", "FLASHINFER") - ): - logger.warning_once( - "%s has interleaved attention, which is currently not " - "supported by the %s backend. Disabling sliding window and " - "capping the max length to the sliding window size (%d).", - self.hf_text_config.model_type, - backend, - self.hf_text_config.sliding_window, - ) - self.disable_sliding_window = True - self.original_max_model_len = self.max_model_len self.max_model_len = self.get_and_verify_max_len(self.max_model_len) # Init multimodal config if needed @@ -359,12 +330,14 @@ class OmniModelConfig(ModelConfig): mm_config_kwargs = dict( limit_per_prompt=limit_mm_per_prompt, + enable_mm_embeds=enable_mm_embeds, media_io_kwargs=media_io_kwargs, mm_processor_kwargs=mm_processor_kwargs, mm_processor_cache_gb=mm_processor_cache_gb, mm_processor_cache_type=mm_processor_cache_type, mm_shm_cache_max_object_size_mb=mm_shm_cache_max_object_size_mb, mm_encoder_tp_mode=mm_encoder_tp_mode, + mm_encoder_attn_backend=mm_encoder_attn_backend, interleave_mm_strings=interleave_mm_strings, skip_mm_profiling=skip_mm_profiling, video_pruning_rate=video_pruning_rate, @@ -374,12 +347,20 @@ class OmniModelConfig(ModelConfig): self.multimodal_config = MultiModalConfig(**mm_config_kwargs) + # Multimodal GGUF models must use original repo for mm processing + if is_gguf(self.tokenizer) and self.is_multimodal_model: + raise ValueError( + "Loading a multimodal GGUF model needs to use original " + "tokenizer. Please specify the unquantized hf model's " + "repo name or path using the --tokenizer argument." + ) + if self.disable_sliding_window: + # Set after get_and_verify_max_len to ensure that max_model_len + # can be correctly capped to sliding window size self.hf_text_config.sliding_window = None - if not self.skip_tokenizer_init: - self._verify_tokenizer_mode() - + # Avoid running try_verify_and_update_config multiple times self.config_updated = False self._verify_quantization() diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 9d69bfeba..7918e1687 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -1,7 +1,10 @@ from __future__ import annotations from collections import defaultdict +from time import time +from vllm.distributed.kv_events import KVEventBatch +from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.logger import init_logger from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler @@ -63,14 +66,6 @@ class OmniARScheduler(VLLMScheduler): scheduler_output: SchedulerOutput, model_runner_output: ModelRunnerOutput, ) -> dict[int, EngineCoreOutputs]: - """Override update_from_output to fix check_stop bug when output_token_ids is empty. - - The original vLLM implementation calls check_stop when pooler_outputs exist, - but doesn't check if output_token_ids is empty. This causes IndexError during - prefill phase when pooler_outputs exist but output_token_ids is still empty. - - This override adds a safety check before calling check_stop. - """ sampled_token_ids = model_runner_output.sampled_token_ids logprobs = model_runner_output.logprobs prompt_logprobs_dict = model_runner_output.prompt_logprobs_dict @@ -81,7 +76,20 @@ class OmniARScheduler(VLLMScheduler): outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None - kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) + if kv_connector_stats and self.connector: + kv_stats = self.connector.get_kv_connector_stats() + if kv_stats: + kv_connector_stats = kv_connector_stats.aggregate(kv_stats) + + failed_kv_load_req_ids = None + if kv_connector_output and kv_connector_output.invalid_block_ids: + # These blocks contain externally computed tokens that failed to + # load. Identify affected requests and adjust their computed token + # count to trigger recomputation of the invalid blocks. + failed_kv_load_req_ids = self._handle_invalid_blocks(kv_connector_output.invalid_block_ids) # 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 @@ -90,6 +98,9 @@ class OmniARScheduler(VLLMScheduler): stopped_preempted_reqs: set[Request] = set() for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): assert num_tokens_scheduled > 0 + if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: + # Skip requests that were recovered from KV load failure + continue request = self.requests.get(req_id) if request is None: # The request is already finished. This can happen if the @@ -110,9 +121,16 @@ class OmniARScheduler(VLLMScheduler): # tokens and rejections. If some tokens are rejected, # num_computed_tokens is decreased by the number of rejected # tokens. - request.num_computed_tokens -= num_rejected + if request.num_computed_tokens > 0: + request.num_computed_tokens -= num_rejected + # If async scheduling, num_output_placeholders also includes + # the scheduled spec tokens count and so is similarly adjusted. + if request.num_output_placeholders > 0: + request.num_output_placeholders -= num_rejected spec_decoding_stats = self.make_spec_decoding_stats( - spec_decoding_stats, num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted + spec_decoding_stats, + num_draft_tokens=num_draft_tokens, + num_accepted_tokens=num_accepted, ) stopped = False @@ -125,28 +143,12 @@ class OmniARScheduler(VLLMScheduler): if new_token_ids: new_token_ids, stopped = self._update_request_with_output(request, new_token_ids) - """ # Stop checking for pooler models. pooler_output = None if pooler_outputs: pooler_output = pooler_outputs[req_index] - stopped = check_stop(request, self.max_model_len, - pooler_output) - """ - - # Stop checking for pooler models. - # FIX: Check if output_token_ids is empty before calling check_stop - # to prevent IndexError during prefill phase when pooler_outputs exist - # but output_token_ids is still empty. - pooler_output = None - if pooler_outputs: - pooler_output = pooler_outputs[req_index] - # Only call check_stop if output_token_ids is not empty - # This prevents IndexError when accessing output_token_ids[-1] if request.output_token_ids: stopped = check_stop(request, self.max_model_len, pooler_output) - # If output_token_ids is empty, we don't call check_stop - # and keep stopped=False (don't stop during prefill phase) if stopped: kv_transfer_params = self._free_request(request) @@ -157,17 +159,13 @@ class OmniARScheduler(VLLMScheduler): # Extract sample logprobs if needed. if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs: - # NOTE: once we support N tokens per step (spec decode), - # the outer lists can be of length > 1. - new_logprobs = logprobs.slice(req_index, req_index + 1) + new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) if new_token_ids and self.structured_output_manager.should_advance(request): - # NOTE: structured_output_request - # should not be None if use_structured_output, we have - # checked above, so safe to ignore type warning - request.structured_output_request.grammar.accept_tokens( # type: ignore[union-attr] - req_id, new_token_ids - ) + struct_output_request = request.structured_output_request + assert struct_output_request is not None + assert struct_output_request.grammar is not None + struct_output_request.grammar.accept_tokens(req_id, new_token_ids) if num_nans_in_logits is not None and req_id in num_nans_in_logits: request.num_nans_in_logits = num_nans_in_logits[req_id] @@ -189,6 +187,7 @@ class OmniARScheduler(VLLMScheduler): kv_transfer_params=kv_transfer_params, trace_headers=request.trace_headers, num_cached_tokens=request.num_cached_tokens, + num_nans_in_logits=request.num_nans_in_logits, ) ) else: @@ -203,8 +202,25 @@ class OmniARScheduler(VLLMScheduler): self.waiting.remove_requests(stopped_preempted_reqs) # KV Connector: update state for finished KV Transfers. - if model_runner_output.kv_connector_output: - self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output) + if kv_connector_output: + self._update_from_kv_xfer_finished(kv_connector_output) + + # collect KV cache events from KV cache manager + events = self.kv_cache_manager.take_events() + + # collect KV cache events from connector + if self.connector is not None: + connector_events = self.connector.take_events() + if connector_events: + if events is None: + events = list(connector_events) + else: + events.extend(connector_events) + + # publish collected KV cache events + if events: + batch = KVEventBatch(ts=time.time(), events=events) + self.kv_event_publisher.publish(batch) # Create EngineCoreOutputs for all clients that have requests with # outputs in this step. diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 430154a55..7acef3bdc 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -1,17 +1,15 @@ import time from collections import defaultdict +from vllm.distributed.kv_events import KVEventBatch from vllm.v1.core.kv_cache_manager import KVCacheBlocks +from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.request_queue import create_request_queue -from vllm.v1.core.sched.scheduler import ( - Request, - RequestStatus, - SchedulerOutput, - SpecDecodingStats, -) from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.core.sched.utils import remove_all from vllm.v1.engine import EngineCoreEventType, EngineCoreOutput, EngineCoreOutputs +from vllm.v1.request import Request, RequestStatus +from vllm.v1.spec_decode.metrics import SpecDecodingStats from vllm_omni.core.sched.output import OmniNewRequestData from vllm_omni.outputs import OmniModelRunnerOutput @@ -26,15 +24,10 @@ class OmniGenerationScheduler(VLLMScheduler): default vLLM scheduling. """ - # Select requests with zero prompt and using pooling - # (diffusion results returned via pooler_output) token_budget = self.max_num_scheduled_tokens - capacity = self.max_num_running_reqs - len(self.running) scheduled_timestamp = time.monotonic() scheduled_new_reqs: list[Request] = [] - scheduled_resumed_reqs: list[Request] = [] - scheduled_running_reqs: list[Request] = [] req_to_new_blocks: dict[str, KVCacheBlocks] = {} num_scheduled_tokens: dict[str, int] = {} @@ -46,7 +39,7 @@ class OmniGenerationScheduler(VLLMScheduler): # Fast path selection and scheduling (treat all as diffusion requests, # independent of pooling_params) - while self.waiting and token_budget > 0 and capacity > 0: + while self.waiting and token_budget > 0 and len(self.running) < self.max_num_running_reqs: request = self.waiting.peek_request() # Uniformly treat as diffusion. A feature flag can be added later # via config or request tag. @@ -81,7 +74,6 @@ class OmniGenerationScheduler(VLLMScheduler): req_to_new_blocks[request.request_id] = new_blocks num_scheduled_tokens[request.request_id] = num_new_tokens token_budget -= num_new_tokens - capacity -= 1 scheduled_new_reqs.append(request) # Return skipped waiting requests @@ -96,25 +88,31 @@ class OmniGenerationScheduler(VLLMScheduler): 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.request_id) - # Assemble SchedulerOutput - new_reqs_data = [ - OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids()) - for req in scheduled_new_reqs - ] + # Assemble SchedulerOutput (align with v0.12.0) + if self.use_v2_model_runner: + # No resumed reqs in fast path; pass prefill_token_ids for new reqs. + new_reqs_data = [ + OmniNewRequestData.from_request( + req, + req_to_new_blocks[req.request_id].get_block_ids(), + getattr(req, "_all_token_ids", None), + ) + for req in scheduled_new_reqs + ] + else: + new_reqs_data = [ + OmniNewRequestData.from_request(req, req_to_new_blocks[req.request_id].get_block_ids()) + for req in scheduled_new_reqs + ] + # No running/resumed reqs scheduled in our fast path cached_reqs_data = self._make_cached_request_data( - scheduled_running_reqs, - scheduled_resumed_reqs, - num_scheduled_tokens, - scheduled_spec_decode_tokens, - req_to_new_blocks, - ) - scheduled_requests = scheduled_new_reqs + scheduled_running_reqs + scheduled_resumed_reqs - structured_output_request_ids, grammar_bitmask = self.get_grammar_bitmask( - scheduled_requests, scheduled_spec_decode_tokens + running_reqs=[], + resumed_reqs=[], + num_scheduled_tokens=num_scheduled_tokens, + spec_decode_tokens=scheduled_spec_decode_tokens, + req_to_new_blocks=req_to_new_blocks, ) total_num_scheduled_tokens = sum(num_scheduled_tokens.values()) @@ -128,14 +126,21 @@ class OmniGenerationScheduler(VLLMScheduler): num_common_prefix_blocks=num_common_prefix_blocks, finished_req_ids=self.finished_req_ids, free_encoder_mm_hashes=self.encoder_cache_manager.get_freed_mm_hashes(), - structured_output_request_ids=structured_output_request_ids, - grammar_bitmask=grammar_bitmask, + preempted_req_ids=set(), ) + # Record the request ids scheduled in this step (v0.12.0 behavior). + self.prev_step_scheduled_req_ids.clear() + self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) + # KVTransfer: package metadata if self.connector is not None: meta = self.connector.build_connector_meta(scheduler_output) scheduler_output.kv_connector_metadata = meta + # EC Connector: package metadata + if self.ec_connector is not None: + ec_meta = self.ec_connector.build_connector_meta(scheduler_output) + scheduler_output.ec_connector_metadata = ec_meta # Update internal state (advance num_computed_tokens, free encoder inputs, # etc.) @@ -171,6 +176,15 @@ class OmniGenerationScheduler(VLLMScheduler): outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None kv_connector_stats = kv_connector_output.kv_connector_stats if kv_connector_output else None + # Merge connector-side stats (align with v0.12.0) + if kv_connector_stats and self.connector: + kv_stats = self.connector.get_kv_connector_stats() + if kv_stats: + kv_connector_stats = kv_connector_stats.aggregate(kv_stats) + + failed_kv_load_req_ids = None + if kv_connector_output and getattr(kv_connector_output, "invalid_block_ids", None): + failed_kv_load_req_ids = self._handle_invalid_blocks(kv_connector_output.invalid_block_ids) # 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 @@ -179,6 +193,9 @@ class OmniGenerationScheduler(VLLMScheduler): stopped_preempted_reqs: set[Request] = set() for req_id, num_tokens_scheduled in num_scheduled_tokens.items(): assert num_tokens_scheduled > 0 + if failed_kv_load_req_ids and req_id in failed_kv_load_req_ids: + # Skip requests that were recovered from KV load failure + continue request = self.requests.get(req_id) if request is None: # The request is already finished. This can happen if the @@ -199,7 +216,8 @@ class OmniGenerationScheduler(VLLMScheduler): # tokens and rejections. If some tokens are rejected, # num_computed_tokens is decreased by the number of rejected # tokens. - request.num_computed_tokens -= num_rejected + if request.num_computed_tokens > 0: + request.num_computed_tokens -= num_rejected spec_decoding_stats = self.make_spec_decoding_stats( spec_decoding_stats, num_draft_tokens=num_draft_tokens, @@ -227,9 +245,7 @@ class OmniGenerationScheduler(VLLMScheduler): # Extract sample logprobs if needed. if request.sampling_params is not None and request.sampling_params.logprobs is not None and logprobs: - # NOTE: once we support N tokens per step (spec decode), - # the outer lists can be of length > 1. - new_logprobs = logprobs.slice(req_index, req_index + 1) + new_logprobs = logprobs.slice_request(req_index, len(new_token_ids)) if new_token_ids and self.structured_output_manager.should_advance(request): # NOTE: structured_output_request should not be None if @@ -245,12 +261,6 @@ class OmniGenerationScheduler(VLLMScheduler): # Get prompt logprobs for this request. prompt_logprobs_tensors = prompt_logprobs_dict.get(req_id) - # Convert pooler_output tensor to dict format expected by OmniEngineCoreOutput - pooling_output_dict = None - if pooler_output is not None: - # Wrap tensor in dict to match OmniEngineCoreOutput.pooling_output type - # which expects Optional[dict[str, torch.Tensor]] - pooling_output_dict = {"model_outputs": pooler_output} if new_token_ids or pooler_output is not None or kv_transfer_params: # Add EngineCoreOutput for this Request. outputs[request.client_index].append( @@ -260,12 +270,13 @@ class OmniGenerationScheduler(VLLMScheduler): finish_reason=request.get_finished_reason(), new_logprobs=new_logprobs, new_prompt_logprobs_tensors=prompt_logprobs_tensors, - pooling_output=pooling_output_dict, + pooling_output=pooler_output, stop_reason=request.stop_reason, events=request.take_events(), kv_transfer_params=kv_transfer_params, trace_headers=request.trace_headers, num_cached_tokens=request.num_cached_tokens, + num_nans_in_logits=request.num_nans_in_logits, ) ) else: @@ -280,8 +291,21 @@ class OmniGenerationScheduler(VLLMScheduler): self.waiting.remove_requests(stopped_preempted_reqs) # KV Connector: update state for finished KV Transfers. - if model_runner_output.kv_connector_output: - self._update_from_kv_xfer_finished(model_runner_output.kv_connector_output) + if kv_connector_output: + self._update_from_kv_xfer_finished(kv_connector_output) + + # Collect and publish KV cache events (align with v0.12.0) + events = self.kv_cache_manager.take_events() + if self.connector is not None: + connector_events = self.connector.take_events() + if connector_events: + if events is None: + events = list(connector_events) + else: + events.extend(connector_events) + if events: + batch = KVEventBatch(ts=time.time(), events=events) + self.kv_event_publisher.publish(batch) # Create EngineCoreOutputs for all clients that have requests with # outputs in this step. diff --git a/vllm_omni/core/sched/output.py b/vllm_omni/core/sched/output.py index bbb38d099..205513ac9 100644 --- a/vllm_omni/core/sched/output.py +++ b/vllm_omni/core/sched/output.py @@ -29,6 +29,7 @@ class OmniNewRequestData(NewRequestData): cls, request: Request, block_ids: tuple[list[int], ...], + prefill_token_ids: list[int] | None = None, ) -> "OmniNewRequestData": """Create OmniNewRequestData from a Request object. @@ -49,5 +50,6 @@ class OmniNewRequestData(NewRequestData): num_computed_tokens=request.num_computed_tokens, lora_request=request.lora_request, prompt_embeds=request.prompt_embeds, + prefill_token_ids=prefill_token_ids, additional_information=request.additional_information, ) diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 460ca9fbc..0398bfbc2 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -13,7 +13,6 @@ from torch import nn from vllm.config import ModelConfig from vllm.config.load import LoadConfig from vllm.logger import init_logger -from vllm.model_executor.model_loader.utils import set_default_torch_dtype from vllm.model_executor.model_loader.weight_utils import ( download_safetensors_index_file_from_hf, download_weights_from_hf, @@ -22,6 +21,7 @@ from vllm.model_executor.model_loader.weight_utils import ( maybe_download_from_modelscope, safetensors_weights_iterator, ) +from vllm.utils.torch_utils import set_default_torch_dtype from vllm_omni.diffusion.data import OmniDiffusionConfig from vllm_omni.diffusion.registry import initialize_model diff --git a/vllm_omni/diffusion/worker/gpu_worker.py b/vllm_omni/diffusion/worker/gpu_worker.py index 40f0b7b30..e7fd05f70 100644 --- a/vllm_omni/diffusion/worker/gpu_worker.py +++ b/vllm_omni/diffusion/worker/gpu_worker.py @@ -9,7 +9,7 @@ import zmq from vllm.config import LoadConfig, VllmConfig, set_current_vllm_config from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.logger import init_logger -from vllm.utils import DeviceMemoryProfiler, GiB_bytes +from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes from vllm_omni.diffusion.cache.selector import get_cache_backend from vllm_omni.diffusion.data import ( diff --git a/vllm_omni/diffusion/worker/npu/npu_worker.py b/vllm_omni/diffusion/worker/npu/npu_worker.py index 1bdb076e5..65f83c504 100644 --- a/vllm_omni/diffusion/worker/npu/npu_worker.py +++ b/vllm_omni/diffusion/worker/npu/npu_worker.py @@ -14,7 +14,7 @@ from vllm.distributed.parallel_state import ( initialize_model_parallel, ) from vllm.logger import init_logger -from vllm.utils import DeviceMemoryProfiler, GiB_bytes +from vllm.utils.mem_utils import DeviceMemoryProfiler, GiB_bytes from vllm_omni.diffusion.data import ( SHUTDOWN_MESSAGE, diff --git a/vllm_omni/distributed/omni_connectors/utils/initialization.py b/vllm_omni/distributed/omni_connectors/utils/initialization.py index cef91915e..608af2dfc 100644 --- a/vllm_omni/distributed/omni_connectors/utils/initialization.py +++ b/vllm_omni/distributed/omni_connectors/utils/initialization.py @@ -292,7 +292,7 @@ def build_stage_connectors( from .config import ConnectorSpec connectors: dict[tuple[str, str], Any] = {} - # 将字典格式的配置转换为ConnectorSpec对象 + # Convert dictionary-formatted config to ConnectorSpec objects stage_connector_specs = {} for input_key, config in connectors_config.items(): if not input_key.startswith("from_stage_"): @@ -310,7 +310,7 @@ def build_stage_connectors( stage_connector_specs[(str(from_stage), str(stage_id))] = connector_spec try: - # 使用统一的连接器创建逻辑 + # Use unified connector creation logic connectors = create_connectors_from_config(stage_connector_specs) except Exception as exc: # pragma: no cover - defensive logging # Fail fast so the stage does not start with missing connectors. diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 1b3cec6ae..79dc25cf4 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -102,6 +102,10 @@ class OmniEngineCoreOutput( # The number of tokens with prefix cache hits. num_cached_tokens: int = 0 + # The number of NaNs in logits. + # A value greater than 0 indicates that the output is corrupted. + num_nans_in_logits: int = 0 + @property def finished(self) -> bool: return self.finish_reason is not None diff --git a/vllm_omni/engine/processor.py b/vllm_omni/engine/input_processor.py similarity index 95% rename from vllm_omni/engine/processor.py rename to vllm_omni/engine/input_processor.py index 0474f9f4a..688387452 100644 --- a/vllm_omni/engine/processor.py +++ b/vllm_omni/engine/input_processor.py @@ -1,6 +1,6 @@ import time from collections.abc import Mapping -from typing import Any +from typing import Any, cast import torch from vllm.config import VllmConfig @@ -9,14 +9,14 @@ from vllm.inputs.parse import split_enc_dec_inputs from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry -from vllm.multimodal.inputs import MultiModalFeatureSpec +from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalUUIDDict from vllm.multimodal.utils import argsort_mm_positions from vllm.platforms import current_platform from vllm.pooling_params import PoolingParams from vllm.sampling_params import SamplingParams -from vllm.transformers_utils.tokenizer import AnyTokenizer +from vllm.tokenizers import TokenizerLike from vllm.utils import length_from_prompt_token_ids_or_embeds -from vllm.v1.engine.processor import Processor +from vllm.v1.engine.input_processor import InputProcessor from vllm_omni.engine import ( AdditionalInformationEntry, @@ -29,7 +29,7 @@ from vllm_omni.inputs.preprocess import OmniInputPreprocessor logger = init_logger(__name__) -class OmniProcessor(Processor): +class OmniInputProcessor(InputProcessor): """Processor for omni models, handling multimodal inputs and embeddings. Extends the base vLLM Processor with support for processing prompt @@ -75,7 +75,7 @@ class OmniProcessor(Processor): def __init__( self, vllm_config: VllmConfig, - tokenizer: AnyTokenizer, + tokenizer: TokenizerLike, mm_registry: MultiModalRegistry = MULTIMODAL_REGISTRY, ): super().__init__(vllm_config, tokenizer, mm_registry) @@ -127,8 +127,6 @@ class OmniProcessor(Processor): ValueError: If data_parallel_rank is out of range or prompt_embeds has incorrect shape """ - # TODO(woosuk): Support pooling models. - # TODO(woosuk): Support encoder-decoder models. self._validate_lora(lora_request) self._validate_params(params) @@ -158,7 +156,7 @@ class OmniProcessor(Processor): # if provided. self._validate_multi_modal_uuids(prompt) if isinstance(prompt, dict): - mm_uuids = prompt.get("multi_modal_uuids") + mm_uuids = cast(MultiModalUUIDDict | None, prompt.get("multi_modal_uuids")) else: mm_uuids = None @@ -171,11 +169,13 @@ class OmniProcessor(Processor): tokenization_kwargs=tokenization_kwargs, mm_uuids=mm_uuids, ) + current_platform.validate_request( prompt=prompt, params=params, processed_inputs=processed_inputs, ) + eos_token_id = self.input_preprocessor.get_eos_token_id() encoder_inputs, decoder_inputs = split_enc_dec_inputs(processed_inputs) @@ -185,7 +185,7 @@ 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: str | None = None if decoder_inputs["type"] == "embeds" else decoder_inputs.get("prompt") + _prompt_str: str | None = 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 @@ -266,7 +266,7 @@ class OmniProcessor(Processor): entries[key] = entry additional_information_payload = AdditionalInformationPayload(entries=entries) - return prompt_str, OmniEngineCoreRequest( + return OmniEngineCoreRequest( request_id=request_id, prompt_token_ids=prompt_token_ids, mm_features=mm_features, diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 5b317e4fc..cf0409c09 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -6,7 +6,7 @@ import torch from vllm.logger import init_logger from vllm.outputs import PoolingRequestOutput from vllm.sampling_params import RequestOutputKind -from vllm.transformers_utils.tokenizer import AnyTokenizer +from vllm.tokenizers import TokenizerLike from vllm.v1.engine import EngineCoreOutput, EngineCoreRequest, FinishReason from vllm.v1.engine.detokenizer import IncrementalDetokenizer from vllm.v1.engine.logprobs import LogprobsProcessor @@ -40,13 +40,14 @@ class OmniRequestState(RequestState): @classmethod def from_new_request( cls, - tokenizer: AnyTokenizer, + tokenizer: TokenizerLike, request: EngineCoreRequest, prompt: str | None, parent_req: ParentRequest | None, request_index: int, queue: Any | None, log_stats: bool, + stream_interval: int, ) -> "OmniRequestState": if sampling_params := request.sampling_params: if not sampling_params.detokenize: @@ -92,6 +93,7 @@ class OmniRequestState(RequestState): arrival_time=request.arrival_time, queue=queue, log_stats=log_stats, + stream_interval=stream_interval, ) def add_multimodal_tensor(self, payload: Any | None, mm_type: str | None) -> None: @@ -196,6 +198,26 @@ class OmniRequestState(RequestState): if not finished and final_only: return None + if self.stream_interval > 1: + assert self.detokenizer is not None + + # Send output request only when + # 1. It has finished, or + # 2. It is the first token, or + # 3. It has reached the stream interval number of tokens + if not ( + finished + or self.sent_tokens_offset == 0 + or len(self.detokenizer.output_token_ids) - self.sent_tokens_offset >= self.stream_interval + ): + return None + + if self.output_kind == RequestOutputKind.DELTA: + # Send tokens from the offset in DELTA mode, otherwise all + # tokens are sent. + new_token_ids = self.detokenizer.output_token_ids[self.sent_tokens_offset :] + self.sent_tokens_offset = len(self.detokenizer.output_token_ids) + request_id = self.request_id output = self._new_completion_output(new_token_ids, finish_reason, stop_reason) @@ -247,7 +269,7 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): def __init__( self, - tokenizer: AnyTokenizer, + tokenizer: TokenizerLike, log_stats: bool, engine_core_output_type: str | None = None, ): @@ -314,9 +336,9 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): request_index=request_index, queue=queue, log_stats=self.log_stats, + stream_interval=self.stream_interval, ) self.request_states[request_id] = req_state - self.lora_states.add_request(req_state) if parent_req: self.parent_requests[parent_req.request_id] = parent_req @@ -422,6 +444,8 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): parent_req = req_state.parent_req if parent_req and not parent_req.child_requests: self.parent_requests.pop(parent_req.request_id, None) + if not self.request_states: + self._requests_drained.set() if not eco.finished: reqs_to_abort.append(req_id) self._update_stats_from_finished(req_state, finish_reason, iteration_stats) diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index ec16d7a8a..f57a90465 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -19,12 +19,13 @@ from vllm.inputs.preprocess import InputPreprocessor from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry +from vllm.plugins.io_processors import get_io_processor from vllm.sampling_params import SamplingParams +from vllm.tokenizers import TokenizerLike, init_tokenizer_from_config from vllm.tracing import init_tracer from vllm.transformers_utils.config import maybe_register_config_serialize_by_value -from vllm.transformers_utils.tokenizer import AnyTokenizer, init_tokenizer_from_configs from vllm.usage.usage_lib import UsageContext -from vllm.utils import Device, deprecate_kwargs +from vllm.utils.func_utils import deprecate_kwargs from vllm.v1.engine.async_llm import AsyncLLM from vllm.v1.engine.core_client import EngineCoreClient from vllm.v1.engine.exceptions import EngineDeadError @@ -44,8 +45,8 @@ from vllm_omni.distributed.ray_utils.utils import ( try_close_ray, ) from vllm_omni.engine.arg_utils import AsyncOmniEngineArgs +from vllm_omni.engine.input_processor import OmniInputProcessor from vllm_omni.engine.output_processor import MultimodalOutputProcessor -from vllm_omni.engine.processor import OmniProcessor from vllm_omni.entrypoints.client_request_state import ClientRequestState from vllm_omni.entrypoints.log_utils import ( OrchestratorMetrics, @@ -114,6 +115,10 @@ class AsyncOmni(EngineClient): self.batch_timeout = cli_args.batch_timeout self._enable_stats: bool = bool(cli_args.log_stats) + # Pause / resume state for async RL workflows. + self._pause_cond = asyncio.Condition() + self._paused = False + base_engine_args = AsyncOmniEngineArgs.from_cli_args(cli_args).__dict__.copy() if cli_args.stage_configs_path is None: @@ -143,6 +148,23 @@ class AsyncOmni(EngineClient): self._stats_file, self._overall_stats_file = init_stats_paths(self._enable_stats, self._log_file) self._initialize_stages(model, cli_args.init_sleep_seconds, cli_args.shm_threshold_bytes, cli_args.init_timeout) + # Here is a duplicated init for the first stage to get the tokenizer, + # input_processor, and io_processor + # TODO: Refactor the code to avoid the duplicated init, we need to pass + # the tokenizer, input_processor, and io_processor back here from the + # first stage. + vllm_config = self.stage_list[0].vllm_config + self.tokenizer = init_tokenizer_from_config(model_config=vllm_config.model_config) + self.input_processor = OmniInputProcessor( + vllm_config=vllm_config, + tokenizer=self.tokenizer, + ) + self.model_config = vllm_config.model_config + self.io_processor = get_io_processor( + vllm_config, + self.model_config.io_processor_plugin, + ) + def _initialize_stages( self, model: str, @@ -304,6 +326,10 @@ class AsyncOmni(EngineClient): Raises: ValueError: If sampling_params_list has incorrect length. """ + # Wait until generation is resumed if the engine is paused. + async with self._pause_cond: + await self._pause_cond.wait_for(lambda: not self._paused) + logger.debug("[Orchestrator] generate() called") # Start output handler on the first call to generate() @@ -634,7 +660,7 @@ class AsyncOmni(EngineClient): async def get_input_preprocessor(self) -> InputPreprocessor: return None - async def get_tokenizer(self) -> AnyTokenizer: + async def get_tokenizer(self) -> TokenizerLike: for stage in self.stage_list: if stage.is_comprehension: return stage.tokenizer @@ -655,7 +681,7 @@ class AsyncOmni(EngineClient): async def reset_mm_cache(self) -> None: pass - async def reset_prefix_cache(self, device: Device | None = None) -> None: + async def reset_prefix_cache(self, reset_running_requests: bool = False) -> bool: pass async def sleep(self, level: int = 1) -> None: @@ -686,6 +712,54 @@ class AsyncOmni(EngineClient): async def stop_profile(self) -> None: raise NotImplementedError("stop_profile() is not implemented for AsyncOmni") + async def pause_generation( + self, + *, + wait_for_inflight_requests: bool = False, + clear_cache: bool = True, + ) -> None: + """ + Pause generation to allow model weight updates. + + New generation/encoding requests are blocked until resume. + + Args: + wait_for_inflight_requests: When ``True`` waits for in-flight + requests to finish before pausing. When ``False`` (default), + immediately aborts any in-flight requests. + clear_cache: Whether to clear KV cache and prefix cache after + draining. Set to ``False`` to preserve cache for faster resume. + Default is ``True`` (clear caches). + """ + + async with self._pause_cond: + if self._paused: + return + self._paused = True + + # Note: AsyncOmni uses a stage-based architecture without a central + # output_processor. For now, we simply set the pause flag and let + # new requests wait. In-flight requests will complete naturally. + # TODO: Implement request abortion for stages if needed. + + # Clear cache if requested + if clear_cache: + await self.reset_prefix_cache() + await self.reset_mm_cache() + + async def resume_generation(self) -> None: + """Resume generation after :meth:`pause_generation`.""" + + async with self._pause_cond: + self._paused = False + self._pause_cond.notify_all() # Wake up all waiting requests + + async def is_paused(self) -> bool: + """Return whether the engine is currently paused.""" + + async with self._pause_cond: + return self._paused + class AsyncOmniStageLLM(AsyncLLM): """Async single-stage LLM engine for use within a stage worker process. @@ -749,14 +823,6 @@ class AsyncOmniStageLLM(AsyncLLM): Returns: None """ - if not envs.VLLM_USE_V1: - raise ValueError( - "Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. " - "This should not happen. As a workaround, try using " - "AsyncLLMEngine.from_vllm_config(...) or explicitly set " - "VLLM_USE_V1=0 or 1 and report this issue on Github." - ) - # Ensure we can serialize custom transformer configs maybe_register_config_serialize_by_value() @@ -773,28 +839,33 @@ class AsyncOmniStageLLM(AsyncLLM): ) if self.model_config.skip_tokenizer_init: - self.tokenizer = None + tokenizer = None else: # Tokenizer (+ ensure liveness if running in another process). - self.tokenizer = init_tokenizer_from_configs(model_config=vllm_config.model_config) + tokenizer = init_tokenizer_from_config(model_config=vllm_config.model_config) - # Processor (converts Inputs --> EngineCoreRequests). - self.processor = OmniProcessor( + # InputProcessor (converts Inputs --> EngineCoreRequests). + self.input_processor = OmniInputProcessor( vllm_config=vllm_config, - tokenizer=self.tokenizer, + tokenizer=tokenizer, mm_registry=mm_registry, ) # OutputProcessor (converts EngineCoreOutputs --> RequestOutput). self.output_processor = MultimodalOutputProcessor( - tokenizer=self.tokenizer, + tokenizer=tokenizer, log_stats=self.log_stats, 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) self.output_processor.tracer = tracer + # Pause / resume state for async RL workflows. + self._pause_cond = asyncio.Condition() + self._paused = False + # EngineCore (starts the engine in background process). self.engine_core = EngineCoreClient.make_async_mp_client( vllm_config=vllm_config, @@ -862,14 +933,6 @@ class AsyncOmniStageLLM(AsyncLLM): client_index: int = 0, disable_log_requests: bool = True, # Deprecated, will be removed ) -> "AsyncLLM": - if not envs.VLLM_USE_V1: - raise ValueError( - "Using V1 AsyncLLMEngine, but envs.VLLM_USE_V1=False. " - "This should not happen. As a workaround, try using " - "AsyncLLMEngine.from_vllm_config(...) or explicitly set " - "VLLM_USE_V1=0 or 1 and report this issue on Github." - ) - # Create the LLMEngine. return cls( vllm_config=vllm_config, diff --git a/vllm_omni/entrypoints/cli/main.py b/vllm_omni/entrypoints/cli/main.py index 9647484f8..6a65d9d6c 100644 --- a/vllm_omni/entrypoints/cli/main.py +++ b/vllm_omni/entrypoints/cli/main.py @@ -16,7 +16,7 @@ def main(): return else: from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup - from vllm.utils import FlexibleArgumentParser + from vllm.utils.argparse_utils import FlexibleArgumentParser import vllm_omni.entrypoints.cli.serve diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 40958bd5e..0485eab97 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -12,7 +12,7 @@ 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.utils.argparse_utils import FlexibleArgumentParser from vllm_omni.entrypoints.openai.api_server import omni_run_server diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index de67bd618..327b88034 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -17,7 +17,7 @@ from vllm.logger import init_logger from vllm.plugins.io_processors import get_io_processor from vllm.sampling_params import SamplingParams from vllm.usage.usage_lib import UsageContext -from vllm.utils import Counter +from vllm.utils.counter import Counter from vllm.v1.engine.llm_engine import LLMEngine from vllm_omni.distributed.omni_connectors import ( @@ -33,8 +33,8 @@ from vllm_omni.distributed.ray_utils.utils import ( try_close_ray, ) from vllm_omni.engine.arg_utils import OmniEngineArgs +from vllm_omni.engine.input_processor import OmniInputProcessor from vllm_omni.engine.output_processor import MultimodalOutputProcessor -from vllm_omni.engine.processor import OmniProcessor from vllm_omni.entrypoints.log_utils import ( OrchestratorMetrics, configure_orchestrator_logger, @@ -624,7 +624,7 @@ class OmniStageLLM(LLM): log_stats=self.llm_engine.log_stats, engine_core_output_type=engine_args.engine_output_type, ) - self.llm_engine.processor = OmniProcessor( + self.llm_engine.input_processor = OmniInputProcessor( vllm_config=self.llm_engine.vllm_config, tokenizer=self.llm_engine.tokenizer ) self.engine_class = type(self.llm_engine) @@ -641,3 +641,5 @@ 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.model_config = self.llm_engine.model_config + self.input_processor = self.llm_engine.input_processor diff --git a/vllm_omni/entrypoints/omni_stage.py b/vllm_omni/entrypoints/omni_stage.py index 4df7c2cbb..8b2050be6 100644 --- a/vllm_omni/entrypoints/omni_stage.py +++ b/vllm_omni/entrypoints/omni_stage.py @@ -22,7 +22,7 @@ from vllm.inputs import TextPrompt from vllm.inputs.preprocess import InputPreprocessor from vllm.logger import init_logger from vllm.sampling_params import SamplingParams -from vllm.transformers_utils.tokenizer import AnyTokenizer +from vllm.tokenizers import TokenizerLike from vllm.usage.usage_lib import UsageContext from vllm.v1.engine import EngineCoreOutput from vllm.v1.engine.async_llm import AsyncLLM @@ -112,7 +112,7 @@ class OmniStage: """ self.vllm_config = vllm_config - def set_tokenizer(self, tokenizer: AnyTokenizer) -> None: + def set_tokenizer(self, tokenizer: TokenizerLike) -> None: """Set the tokenizer for this stage. Args: diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 000e95a81..ca30bb901 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -37,8 +37,8 @@ from vllm.entrypoints.openai.tool_parsers import ToolParserManager from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer 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.tokenizers import MistralTokenizer +from vllm.utils.system_utils import decorate_logs from vllm_omni.diffusion.data import DiffusionParallelConfig, OmniDiffusionConfig from vllm_omni.diffusion.utils.hf_utils import is_diffusion_model @@ -305,13 +305,8 @@ async def build_async_omni_from_stage_config( """ # V1 AsyncLLM. - assert envs.VLLM_USE_V1 - if disable_frontend_multiprocessing: - logger.warning( - "V1 is enabled, but got --disable-frontend-multiprocessing. " - "To disable frontend multiprocessing, set VLLM_USE_V1=0." - ) + logger.warning("V1 is enabled, but got --disable-frontend-multiprocessing.") async_omni: EngineClient | None = None @@ -358,7 +353,7 @@ async def omni_init_app_state( state.engine_client = engine_client state.log_stats = not args.disable_log_stats state.vllm_config = vllm_config - model_config = vllm_config.model_config + _model_config = vllm_config.model_config state.log_stats = not args.disable_log_stats # For omni models @@ -416,14 +411,12 @@ async def omni_init_app_state( state.openai_serving_models = OpenAIServingModels( engine_client=engine_client, - model_config=model_config, base_model_paths=base_model_paths, lora_modules=lora_modules, ) await state.openai_serving_models.init_static_loras() state.openai_serving_chat = OmniOpenAIServingChat( engine_client, - model_config, state.openai_serving_models, args.response_role, request_logger=request_logger, diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 5d42ae3d2..f079be4e2 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -61,13 +61,14 @@ from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.outputs import RequestOutput from vllm.sampling_params import SamplingParams -from vllm.transformers_utils.tokenizer import AnyTokenizer, MistralTokenizer -from vllm.transformers_utils.tokenizers import ( +from vllm.tokenizers import TokenizerLike +from vllm.tokenizers.mistral import ( + MistralTokenizer, maybe_serialize_tool_calls, truncate_tool_call_ids, validate_request_params, ) -from vllm.utils import as_list +from vllm.utils.collection_utils import as_list from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures from vllm_omni.outputs import OmniRequestOutput @@ -290,7 +291,7 @@ class OmniOpenAIServingChat(OpenAIServingChat): async def _preprocess_chat( self, request: ChatLikeRequest | ResponsesRequest, - tokenizer: AnyTokenizer, + tokenizer: TokenizerLike, messages: list[ChatCompletionMessageParam], chat_template: str | None, chat_template_content_format: ChatTemplateContentFormatOption, @@ -299,7 +300,7 @@ class OmniOpenAIServingChat(OpenAIServingChat): tool_dicts: list[dict[str, Any]] | None = None, documents: list[dict[str, str]] | None = None, chat_template_kwargs: dict[str, Any] | None = None, - tool_parser: Callable[[AnyTokenizer], ToolParser] | None = None, + tool_parser: Callable[[TokenizerLike], ToolParser] | None = None, add_special_tokens: bool = False, ) -> tuple[ list[ConversationMessage], @@ -450,7 +451,7 @@ class OmniOpenAIServingChat(OpenAIServingChat): request_id: str, model_name: str, conversation: list[ConversationMessage], - tokenizer: AnyTokenizer, + tokenizer: TokenizerLike, request_metadata: RequestResponseMetadata, ) -> ErrorResponse | ChatCompletionResponse: created_time = int(time.time()) @@ -542,7 +543,7 @@ class OmniOpenAIServingChat(OpenAIServingChat): self, request: ChatCompletionRequest, omni_outputs: OmniRequestOutput, - tokenizer: AnyTokenizer, + tokenizer: TokenizerLike, conversation: list[ConversationMessage], role: str, ): diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 262b18bf7..4c1ee5388 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -36,7 +36,7 @@ class OmniInputPreprocessor(InputPreprocessor): inputs = self._process_multimodal( prompt_token_ids, multi_modal_data, - parsed_content.get("mm_processor_kwargs"), + parsed_content.get("mm_processor_kwargs") or {}, tokenization_kwargs=tokenization_kwargs, mm_uuids=mm_uuids, ) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py index 331594f95..99a80dc82 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni.py @@ -15,6 +15,12 @@ from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP +from vllm.model_executor.models.qwen2_5_omni_thinker import ( + Qwen2_5OmniConditionalGenerationMixin, + Qwen2_5OmniThinkerDummyInputsBuilder, + Qwen2_5OmniThinkerMultiModalProcessor, + Qwen2_5OmniThinkerProcessingInfo, +) from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix # from vllm.model_executor.models.qwen2_code2wav_dit import Qwen2Code2wav @@ -26,12 +32,6 @@ from vllm.v1.sample.sampler import Sampler from vllm_omni.model_executor.model_loader.weight_utils import download_weights_from_hf_specific from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( - Qwen2_5OmniConditionalGenerationMixin, - Qwen2_5OmniThinkerDummyInputsBuilder, - Qwen2_5OmniThinkerMultiModalProcessor, - Qwen2_5OmniThinkerProcessingInfo, -) from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights, split_list_into_ranges from vllm_omni.model_executor.models.vision import get_llm_pos_ids_for_vision from vllm_omni.utils.platform_utils import is_npu @@ -161,18 +161,21 @@ class Qwen2_5OmniForConditionalGeneration( return self.model.sampler return Sampler() - def get_input_embeddings( + def embed_input_ids( self, input_ids: torch.Tensor, multimodal_embeddings=None, + is_multimodal=None, ) -> torch.Tensor: if self.model_stage == "code2wav": return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size()) - return self.model.get_input_embeddings(input_ids, multimodal_embeddings) + return self.model.embed_input_ids( + input_ids=input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal + ) - def get_multimodal_embeddings(self, **kwargs): + def embed_multimodal(self, **kwargs): # Delegate to thinker model for multimodal processing - return self.model.get_multimodal_embeddings(**kwargs) + return self.model.embed_multimodal(**kwargs) def last_index_of(self, list, value): return len(list) - 1 - list[::-1].index(value) @@ -258,7 +261,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.reshape(-1, text_hidden_states.shape[-1])), multimodal_outputs=None, ) @@ -288,7 +291,7 @@ class Qwen2_5OmniForConditionalGeneration( # Ensure we have base embeddings when only ids are provided if inputs_embeds is None and input_ids is not None: - inputs_embeds = self.talker.get_input_embeddings(input_ids) + inputs_embeds = self.talker.embed_input_ids(input_ids) # ------- Request-scoped additional information (no cross-request concat) ------- request_ids: list[str] | None = kwargs.get("request_ids") # ordered @@ -456,7 +459,7 @@ class Qwen2_5OmniForConditionalGeneration( dtype=torch.bfloat16, ).to(self._module_device(self.model)), self.talker.thinker_to_talker_proj( - self.talker.get_input_embeddings( + self.talker.embed_input_ids( torch.tensor([TALKER_CODEC_BOS_TOKEN_ID, TALKER_CODEC_EOS_TOKEN_ID]) .to(torch.bfloat16) .to(self._module_device(self.model)) @@ -790,7 +793,7 @@ class Qwen2_5OmniForConditionalGeneration( output_prompt_embeds, output_token_ids, ): - processed_output_token_embeds = output_prompt_embeds + self.talker.get_input_embeddings( + processed_output_token_embeds = output_prompt_embeds + self.talker.embed_input_ids( output_token_ids ) # for decode return output_token_ids, processed_output_token_embeds diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py index 1a0edbed9..e784146b4 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_talker.py @@ -12,7 +12,6 @@ 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.qwen2_5_omni_thinker import ( - Qwen2_5OmniConditionalGenerationMixin, Qwen2_5OmniThinkerDummyInputsBuilder, Qwen2_5OmniThinkerMultiModalProcessor, Qwen2_5OmniThinkerProcessingInfo, @@ -23,7 +22,6 @@ from vllm.model_executor.models.utils import ( WeightsMapper, init_vllm_registered_model, maybe_prefix, - merge_multimodal_embeddings, ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors @@ -31,6 +29,8 @@ from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import Qwen2_5OmniConditionalGenerationMixin + @MULTIMODAL_REGISTRY.register_processor( Qwen2_5OmniThinkerMultiModalProcessor, @@ -58,7 +58,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration( super().__init__() config: Qwen2_5OmniTalkerConfig = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config - + self.vllm_config = vllm_config self.prefix = prefix self.quant_config = quant_config @@ -103,26 +103,24 @@ class Qwen2_5OmniTalkerForConditionalGeneration( return Sampler() - def get_input_embeddings( + def embed_input_ids( self, input_ids: torch.Tensor, multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + handle_oov_mm_token: bool = False, ) -> torch.Tensor: - inputs_embeds = self.language_model.get_input_embeddings(input_ids) - if multimodal_embeddings is not None and len(multimodal_embeddings) != 0: - # TODO (ywang96): support overlapping modalitiy embeddings so that - # `use_audio_in_video` will work on V1. - inputs_embeds = merge_multimodal_embeddings( - input_ids, - inputs_embeds, - multimodal_embeddings, - [ - self.config.image_token_index, - self.config.video_token_index, - self.config.audio_token_index, - ], - ) - return inputs_embeds + # This is to satisfy the type checker for each overload + if multimodal_embeddings is None or is_multimodal is None: + return super().embed_input_ids(input_ids) + + return super().embed_input_ids( + input_ids, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + handle_oov_mm_token=handle_oov_mm_token, + ) def forward( self, @@ -139,7 +137,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration( inputs_embeds = None elif inputs_embeds is None: # for profile_run: - inputs_embeds = self.get_input_embeddings(input_ids) + inputs_embeds = self.embed_input_ids(input_ids) input_ids = None @@ -213,7 +211,7 @@ class Qwen2_5OmniTalkerForConditionalGeneration( 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: + def embed_multimodal(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/qwen2_5_omni_thinker.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py index 405c18e74..78737ebf3 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_5_omni_thinker.py @@ -1,694 +1,96 @@ -from collections import defaultdict -from collections.abc import Callable, Iterable, Mapping, Sequence -from functools import partial -from typing import Annotated, Any, Literal +"""Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker (v0.12) with minimal overrides.""" + +from collections.abc import Iterable +from typing import Any import torch -import torch.nn as nn -from transformers.feature_extraction_utils import BatchFeature -from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import Qwen2_5OmniConfig, Qwen2_5OmniThinkerConfig -from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import Qwen2_5OmniAudioEncoder -from transformers.models.qwen2_5_omni.processing_qwen2_5_omni import Qwen2_5OmniProcessor -from transformers.models.whisper import WhisperFeatureExtractor +from torch import nn +from transformers.models.qwen2_5_omni.configuration_qwen2_5_omni import ( + Qwen2_5OmniThinkerConfig, +) +from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import ( + Qwen2_5OmniAudioEncoder, +) from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMRoPE, SupportsMultiModal, SupportsPP +from vllm.model_executor.models.interfaces import ( + MultiModalEmbeddings, + SupportsLoRA, + SupportsMRoPE, + SupportsMultiModal, + SupportsPP, +) +from vllm.model_executor.models.module_mapping import MultiModelKeys +from vllm.model_executor.models.qwen2_5_omni_thinker import ( + Qwen2_5OmniAudioFeatureInputs, + Qwen2_5OmniThinkerDummyInputsBuilder, + Qwen2_5OmniThinkerMultiModalProcessor, + Qwen2_5OmniThinkerProcessingInfo, + get_llm_pos_ids_for_vision, + split_list_into_ranges, +) +from vllm.model_executor.models.qwen2_5_omni_thinker import ( + Qwen2_5OmniConditionalGenerationMixin as Qwen2_5OmniConditionalGenerationMixinBase, +) from vllm.model_executor.models.qwen2_5_vl import ( Qwen2_5_VisionTransformer, Qwen2_5_VLImageEmbeddingInputs, Qwen2_5_VLImageInputs, Qwen2_5_VLImagePixelInputs, - Qwen2_5_VLProcessingInfo, Qwen2_5_VLVideoEmbeddingInputs, Qwen2_5_VLVideoInputs, Qwen2_5_VLVideoPixelInputs, ) -from vllm.model_executor.models.qwen2_audio import ( - Qwen2AudioProcessingInfo, - _get_feat_extract_output_lengths, -) -from vllm.model_executor.models.qwen2_vl import Qwen2VLMultiModalDataParser from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, init_vllm_registered_model, maybe_prefix, - merge_multimodal_embeddings, ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( - ImageItem, - ModalityData, - MultiModalDataDict, - MultiModalFieldConfig, - MultiModalKwargsItems, - NestedTensors, + MultiModalFeatureSpec, ) -from vllm.multimodal.parse import ( - AudioProcessorItems, - DictEmbeddingItems, - ModalityDataItems, - MultiModalDataItems, - MultiModalDataParser, -) -from vllm.multimodal.processing import ( - BaseMultiModalProcessor, - MultiModalPromptUpdates, - PlaceholderFeaturesInfo, - PromptReplacement, - PromptUpdate, - PromptUpdateDetails, - ResolvedPromptUpdate, -) -from vllm.multimodal.profiling import BaseDummyInputsBuilder from vllm.sequence import IntermediateTensors -from vllm.transformers_utils.tokenizer import decode_tokens, encode_tokens -from vllm.utils.tensor_schema import TensorSchema, TensorShape - -from vllm_omni.model_executor.layers.mrope import MRotaryEmbedding try: import flash_attn except (ImportError, ModuleNotFoundError): flash_attn = None - logger = init_logger(__name__) -class Qwen2_5OmniAudioFeatureInputs(TensorSchema): - """ - Dimensions: - - na: Number of audios - - nmb: Number of mel bins - - msl: Maximum sequence length - - tsl: Total sequence length - """ - - type: Literal["audio_features"] - input_features: Annotated[ - torch.Tensor | list[torch.Tensor], - TensorShape("nmb", "tsl"), - ] - - feature_attention_mask: Annotated[ - torch.Tensor, - TensorShape("na", "msl"), - ] - - -def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]): - audio_feature_lengths = hf_inputs.get("audio_feature_lengths", torch.empty((0,))) - - image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3))) - image_grid_sizes = image_grid_thw.prod(-1) - - video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3))) - video_grid_sizes = video_grid_thw.prod(-1) - - return dict( - input_audio_features=MultiModalFieldConfig.flat_from_sizes("audio", audio_feature_lengths, dim=1), - feature_attention_mask=MultiModalFieldConfig.batched("audio"), - audio_feature_lengths=MultiModalFieldConfig.batched("audio"), - pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes), - image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_grid_sizes), - image_grid_thw=MultiModalFieldConfig.batched("image"), - pixel_values_videos=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes), - video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes), - video_grid_thw=MultiModalFieldConfig.batched("video"), - second_per_grid_ts=MultiModalFieldConfig.batched("video"), - ) - - -def create_qwen2_5_omni_thinker_field_factory( - spatial_merge_size: int, -) -> Callable[[Mapping[str, torch.Tensor]], Mapping[str, MultiModalFieldConfig]]: - def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]): - audio_feature_lengths = hf_inputs.get("audio_feature_lengths", torch.empty((0,))) - - image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3))) - image_pixel_grid_sizes = image_grid_thw.prod(-1) - image_embed_grid_sizes = image_pixel_grid_sizes // spatial_merge_size // spatial_merge_size - - video_grid_thw = hf_inputs.get("video_grid_thw", torch.empty((0, 3))) - video_grid_sizes = video_grid_thw.prod(-1) - video_embed_grid_sizes = video_grid_sizes // spatial_merge_size // spatial_merge_size - - num_videos = len(video_grid_sizes) - - return dict( - input_audio_features=MultiModalFieldConfig.flat_from_sizes("audio", audio_feature_lengths, dim=1), - feature_attention_mask=MultiModalFieldConfig.batched("audio"), - audio_feature_lengths=MultiModalFieldConfig.batched("audio"), - pixel_values=MultiModalFieldConfig.flat_from_sizes("image", image_pixel_grid_sizes), - image_embeds=MultiModalFieldConfig.flat_from_sizes("image", image_embed_grid_sizes), - image_grid_thw=MultiModalFieldConfig.batched("image"), - pixel_values_videos=MultiModalFieldConfig.flat_from_sizes("video", video_grid_sizes), - video_embeds=MultiModalFieldConfig.flat_from_sizes("video", video_embed_grid_sizes), - video_grid_thw=MultiModalFieldConfig.batched("video"), - second_per_grid_ts=MultiModalFieldConfig.batched("video"), - use_audio_in_video=MultiModalFieldConfig.shared("video", num_videos), - ) - - return _qwen2_5_omni_thinker_field_config - - -class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser): - def _parse_audio_data( - self, - data: dict[str, torch.Tensor] | ModalityData[ImageItem], - ) -> ModalityDataItems[Any, Any]: - if isinstance(data, dict): - return DictEmbeddingItems( - data, - modality="audio", - required_fields={"input_audio_features", "audio_feature_lengths"}, - fields_factory=_qwen2_5_omni_thinker_field_config, - ) - - return super()._parse_audio_data(data) - - -class Qwen2_5OmniThinkerProcessingInfo(Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo): - def get_hf_config(self): - return self.ctx.get_hf_config(Qwen2_5OmniConfig).thinker_config - - def get_hf_processor(self, **kwargs: object) -> Qwen2_5OmniProcessor: - return self.ctx.get_hf_processor( - Qwen2_5OmniProcessor, - use_fast=kwargs.pop("use_fast", True), - **kwargs, - ) - - def get_feature_extractor(self, **kwargs: object): - hf_processor = self.get_hf_processor(**kwargs) - feature_extractor = hf_processor.feature_extractor # type: ignore - assert isinstance(feature_extractor, WhisperFeatureExtractor) - return feature_extractor - - def get_supported_mm_limits(self) -> Mapping[str, int | None]: - return {"audio": None, "image": None, "video": None} - - -class Qwen2_5OmniThinkerDummyInputsBuilder(BaseDummyInputsBuilder[Qwen2_5OmniThinkerProcessingInfo]): - def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str: - num_audios = mm_counts.get("audio", 0) - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - - hf_processor = self.info.get_hf_processor() - - audio_token: str = hf_processor.audio_token - image_token: str = hf_processor.image_token - video_token: str = hf_processor.video_token - - return audio_token * num_audios + image_token * num_images + video_token * num_videos - - def get_dummy_mm_data( - self, - seq_len: int, - mm_counts: Mapping[str, int], - ) -> MultiModalDataDict: - num_audios = mm_counts.get("audio", 0) - num_images = mm_counts.get("image", 0) - num_videos = mm_counts.get("video", 0) - - feature_extractor = self.info.get_feature_extractor() - - target_audio_length = ( - min( - feature_extractor.chunk_length, - 30, - ) - * feature_extractor.sampling_rate - ) - target_width, target_height = self.info.get_image_size_with_most_features() - target_num_frames = self.info.get_num_frames_with_most_features(seq_len, mm_counts) - - mm_data = { - "audio": self._get_dummy_audios(length=target_audio_length, num_audios=num_audios), - "image": self._get_dummy_images(width=target_width, height=target_height, num_images=num_images), - "video": self._get_dummy_videos( - width=target_width, height=target_height, num_frames=target_num_frames, num_videos=num_videos - ), - } - - return mm_data - - -class Qwen2_5OmniThinkerMultiModalProcessor(BaseMultiModalProcessor[Qwen2_5OmniThinkerProcessingInfo]): - def _get_data_parser(self) -> MultiModalDataParser: - feature_extractor = self.info.get_feature_extractor() - return Qwen2_5OmniThinkerMultiModalDataParser( - spatial_merge_size=self.info.get_hf_config().vision_config.spatial_merge_size, - target_sr=feature_extractor.sampling_rate, - ) - - def _call_hf_processor( - self, - prompt: str, - mm_data: Mapping[str, object], - mm_kwargs: Mapping[str, object], - tok_kwargs: Mapping[str, object], - ) -> BatchFeature: - mm_data = dict(mm_data) - audios = mm_data.pop("audios", []) - - # NOTE: WhisperFeatureExtractor cannot handle empty list of audios - if audios: - # NOTE: Qwen2.5-Omni processor accept "audio" - mm_data["audio"] = audios - mm_kwargs = dict( - **mm_kwargs, - ) - - hf_inputs = super()._call_hf_processor( - prompt=prompt, - mm_data=mm_data, - mm_kwargs=mm_kwargs, - tok_kwargs=tok_kwargs, - ) - - input_features = hf_inputs.pop("input_features", None) - feature_attention_mask = hf_inputs.get("feature_attention_mask", None) - if "input_audio_features" not in hf_inputs and input_features is not None: - if feature_attention_mask is not None: - input_features = input_features.permute(0, 2, 1)[feature_attention_mask.bool()].permute(1, 0) - hf_inputs["input_audio_features"] = input_features - if "audio_feature_lengths" not in hf_inputs and feature_attention_mask is not None: - hf_inputs["audio_feature_lengths"] = feature_attention_mask.sum(-1) - - video_second_per_grid = hf_inputs.get("video_second_per_grid", None) - if video_second_per_grid is not None: - hf_inputs["second_per_grid_ts"] = video_second_per_grid - - use_audio_in_video = mm_kwargs.get("use_audio_in_video", False) - hf_inputs["use_audio_in_video"] = torch.tensor(use_audio_in_video) - - return hf_inputs - - def _get_mm_fields_config( - self, - hf_inputs: BatchFeature, - hf_processor_mm_kwargs: Mapping[str, object], - ) -> Mapping[str, MultiModalFieldConfig]: - return create_qwen2_5_omni_thinker_field_factory(self.info.get_hf_config().vision_config.spatial_merge_size)( - hf_inputs - ) - - def _apply_prompt_updates( - self, - token_ids: list[int], - mm_prompt_updates: MultiModalPromptUpdates, - ) -> tuple[list[int], Mapping[str, list[PlaceholderFeaturesInfo]]]: - tokenizer = self.info.get_tokenizer() - - new_token_ids, match_result = self._apply_token_matches( - token_ids, - mm_prompt_updates, - ) - - # If the search text does not represent a special token, - # it may have different token IDs in the prompt, because - # the tokens may go across the boundaries of the search text. - # ---- - # e.g. when searching for "foo" in "food", if "food" itself makes - # up a token, then the token ID of "foo" will not appear at all - # ---- - # Since it is inefficient to search for all possible tokenizations - # of the search text in the prompt, we instead perform string-based - # updates on the decoded token IDs, then encode them back. - if not all(all(update_idx is not None for update_idx in update_idxs) for update_idxs in match_result.values()): - new_text, match_result = self._apply_text_matches( - decode_tokens(tokenizer, token_ids), - mm_prompt_updates, - ) - - new_token_ids = encode_tokens( - tokenizer, - new_text, - add_special_tokens=False, - ) - - matched_updates = defaultdict[str, list[Sequence[ResolvedPromptUpdate]]](list) - unmatched_audio_items = [] - - for modality, update_idxs in match_result.items(): - for item_idx, update_idx in enumerate(update_idxs): - if update_idx is None: - # Check if this is audio that might be embedded in video - if modality == "audio": - unmatched_audio_items.append(item_idx) - continue - else: - assert False, f"Failed to apply prompt replacement for mm_items[{modality!r}][{item_idx}]" - - matched_updates[modality].append([mm_prompt_updates[modality][item_idx][update_idx]]) - - # If there are unmatched audio items, check if we can derive them from video - if unmatched_audio_items: - # Check if video exists in the original updates (not just matched ones) - if "video" in mm_prompt_updates: - num_videos = len(mm_prompt_updates.get("video", [])) - num_audios = len(mm_prompt_updates.get("audio", [])) - # If counts match, audio might be embedded in video - skip assertion - # The placeholders will be derived from video in _maybe_apply_prompt_updates - if num_audios == num_videos and "video" in matched_updates: - pass # Will be handled by deriving from video placeholders - else: - # Audio items exist but can't be matched - # Check if video was also matched - if "video" not in matched_updates: - # Neither audio nor video matched - this is a real error - assert False, ( - f"Failed to apply prompt replacement for {len(unmatched_audio_items)} " - f"audio item(s). Audio items cannot be matched in prompt." - ) - elif num_audios != num_videos: - # Audio and video counts don't match - assert False, ( - f"Failed to apply prompt replacement for {len(unmatched_audio_items)} " - f"audio item(s). Audio items cannot be matched in prompt and " - f"audio count ({num_audios}) does not match video count ({num_videos})." - ) - else: - # Video matched but audio didn't, and counts match - allow it - pass - else: - # Audio items can't be matched and there are no video items - assert False, ( - f"Failed to apply prompt replacement for {len(unmatched_audio_items)} " - f"audio item(s). Audio items cannot be matched in prompt." - ) - - placeholders = self._find_mm_placeholders( - new_token_ids, - dict(matched_updates), - ) - - return new_token_ids, placeholders - - def _maybe_apply_prompt_updates( - self, - mm_items: MultiModalDataItems, - prompt_ids: list[int], - mm_kwargs: MultiModalKwargsItems, - mm_prompt_updates: MultiModalPromptUpdates, - is_update_applied: bool, - ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]: - """ - Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`. - """ - mm_item_counts = mm_items.get_all_counts() - self._validate_mm_kwargs(mm_kwargs, mm_item_counts) - - use_audio_in_video = False - if "video" in mm_kwargs: - video_items = [item for item in mm_kwargs["video"] if item is not None] - # only check video items (if there are any) - if video_items: - use_audio_in_video = all(item["use_audio_in_video"].data for item in video_items) - - if is_update_applied: - mm_placeholders = self._find_mm_placeholders( - prompt_ids, - mm_prompt_updates, - ) - self._validate_mm_placeholders( - mm_placeholders, - mm_item_counts, - ) - - else: - # When use_audio_in_video=True, audio tokens are not in the prompt, - # so we need to filter out audio updates before applying replacements - if use_audio_in_video and "audio" in mm_prompt_updates: - # Remove audio from prompt updates (it won't match anything) - filtered_updates = {k: v for k, v in mm_prompt_updates.items() if k != "audio"} - prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - filtered_updates, - ) - # Derive audio placeholders from video placeholders - mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates) - else: - # Apply prompt updates normally - # _apply_prompt_updates will handle unmatched audio items - # when video exists with matching counts - prompt_ids, mm_placeholders = self._apply_prompt_updates( - prompt_ids, - mm_prompt_updates, - ) - - # Check if audio placeholders are missing but should exist - # This can happen when audio items are embedded in video tokens - if "audio" in mm_prompt_updates and "audio" not in mm_placeholders and "video" in mm_placeholders: - num_audios = len(mm_prompt_updates.get("audio", [])) - num_videos = len(mm_placeholders.get("video", [])) - # If counts match, derive audio placeholders from video - if num_audios == num_videos: - mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates) - self._validate_mm_placeholders( - mm_placeholders, - mm_item_counts, - ) - - tokenizer = self.info.get_tokenizer() - prompt = decode_tokens(tokenizer, prompt_ids) - - return prompt_ids, prompt, mm_placeholders - - def _get_prompt_updates( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, Any], - out_mm_kwargs: MultiModalKwargsItems, - ) -> Sequence[PromptUpdate]: - processor = self.info.get_hf_processor(**hf_processor_mm_kwargs) - tokenizer = self.info.get_tokenizer() - image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs) - vocab = tokenizer.get_vocab() - - audio_token = processor.audio_token - image_token = processor.image_token - video_token = processor.video_token - audio_token_id = vocab[audio_token] - image_token_id = vocab[image_token] - video_token_id = vocab[video_token] - - out_mm_data = out_mm_kwargs.get_data() - audio_feature_lengths = out_mm_data.get("audio_feature_lengths") - feature_attention_mask = out_mm_data.get("feature_attention_mask") - if audio_feature_lengths is None and feature_attention_mask is None: - audio_output_lengths = [] - elif audio_feature_lengths is not None: - _, audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths) - audio_output_lengths = audio_output_lens.tolist() - elif feature_attention_mask is not None: - assert isinstance(feature_attention_mask, torch.Tensor) - _, audio_output_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1)) - audio_output_lengths = audio_output_lens.tolist() - - # number of audios read from video. - audio_in_video_item_idx = 0 - - def get_replacement_qwen2_audio(item_idx: int): - item_idx += audio_in_video_item_idx - - num_features = audio_output_lengths[item_idx] - if num_features == 0: - audios = mm_items.get_items("audio", AudioProcessorItems) - audio = audios.get(item_idx) - raise ValueError( - f"The audio {audio} (len={len(audio)}) is too short to be represented inside the model" - ) - - return [audio_token_id] * num_features - - def get_replacement_qwen2_vision(item_idx: int, modality: str): - grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx] - assert isinstance(grid_thw, torch.Tensor) - merge_length = image_processor.merge_size**2 - - token_id = image_token_id if modality == "image" else video_token_id - return [token_id] * (int(grid_thw.prod()) // merge_length) - - use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False) - thinker_config = self.info.get_hf_config() - - def get_replacement_qwen2_use_audio_in_video(item_idx: int): - nonlocal audio_in_video_item_idx - - audio_num_features = audio_output_lengths[audio_in_video_item_idx + item_idx] - video_grid_thw = out_mm_data["video_grid_thw"][item_idx] - - audio_in_video_item_idx += 1 - - second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None) - if second_per_grid_ts: - video_second_per_grid_t = second_per_grid_ts[item_idx] - else: - video_second_per_grid_t = 1.0 - - placeholder = MRotaryEmbedding.omni_get_updates_use_audio_in_video( - thinker_config=thinker_config, - audio_len=audio_num_features, - video_grid_thw=video_grid_thw, - video_second_per_grid_t=video_second_per_grid_t, - ) - - return PromptUpdateDetails.select_token_id(placeholder, embed_token_id=video_token_id) - - video_replacement_fn = ( - get_replacement_qwen2_use_audio_in_video - if use_audio_in_video - else partial(get_replacement_qwen2_vision, modality="video") - ) - - return [ - PromptReplacement( - modality="audio", - target=audio_token, - replacement=get_replacement_qwen2_audio, - ), - PromptReplacement( - modality="image", - target=image_token, - replacement=partial(get_replacement_qwen2_vision, modality="image"), - ), - PromptReplacement( - modality="video", - target=video_token, - replacement=video_replacement_fn, - ), - ] - - def _derive_audio_from_video_placeholders( - self, - placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], - mm_prompt_updates: MultiModalPromptUpdates, - ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: - """ - Helper to derive audio placeholders from video placeholders when - use_audio_in_video=True. - """ - if "video" not in placeholders: - return placeholders - - # Validate audio and video counts match - num_videos = len(placeholders["video"]) - num_audios = len(mm_prompt_updates.get("audio", [])) - if num_audios != num_videos: - raise ValueError( - f"use_audio_in_video requires equal number of audio and video items, got {num_audios=}, {num_videos=}" - ) - - tokenizer = self.info.get_tokenizer() - processor = self.info.get_hf_processor() - audio_token_id = tokenizer.get_vocab()[processor.audio_token] - - result_placeholders = dict(placeholders) - audio_placeholders = [] - - # Each video is paired with one audio - for video_idx, video_placeholder in enumerate(placeholders["video"]): - # Create is_embed mask selecting only audio tokens - audio_is_embed = torch.tensor(video_placeholder.tokens) == audio_token_id - - audio_placeholder = PlaceholderFeaturesInfo( - modality="audio", - item_idx=video_idx, - start_idx=video_placeholder.start_idx, - tokens=video_placeholder.tokens, - is_embed=audio_is_embed, - ) - audio_placeholders.append(audio_placeholder) - - result_placeholders["audio"] = audio_placeholders - return result_placeholders - - def _apply_hf_processor_main( - self, - prompt: str | list[int], - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - *, - enable_hf_prompt_update: bool, - ) -> tuple[list[int], BatchFeature, bool]: - """ - Qwen2.5-Omni reimplements this function to handle text only. - """ - if isinstance(prompt, str): - if enable_hf_prompt_update: - return self._apply_hf_processor_text_mm( - prompt_text=prompt, - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - tokenizer = self.info.get_tokenizer() - prompt_ids = encode_tokens(tokenizer, prompt) - else: - prompt_ids = self._apply_hf_processor_tokens_only(prompt) - - mm_processed_data = self._apply_hf_processor_mm_only( - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - return prompt_ids, mm_processed_data, False - - def _apply_hf_processor_mm_only( - self, - mm_items: MultiModalDataItems, - hf_processor_mm_kwargs: Mapping[str, object], - tokenization_kwargs: Mapping[str, object], - ) -> BatchFeature: - """ - Qwen2.5-Omni reimplements this function to handle `use_audio_in_video`. - """ - mm_counts = mm_items.get_all_counts() - - use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False) - if use_audio_in_video and "video" in mm_counts: - assert "audio" in mm_counts - mm_counts["audio"] -= mm_counts["video"] - - _, mm_processed_data, _ = self._apply_hf_processor_text_mm( - prompt_text=self.dummy_inputs.get_dummy_text(mm_counts), - mm_items=mm_items, - hf_processor_mm_kwargs=hf_processor_mm_kwargs, - tokenization_kwargs=tokenization_kwargs, - ) - - return mm_processed_data - - -class Qwen2_5OmniConditionalGenerationMixin: - def _validate_and_reshape_mm_tensor(self, mm_input: object, name: str, dim: int = 0) -> torch.Tensor: - if not isinstance(mm_input, (torch.Tensor, list)): - raise ValueError(f"Incorrect type of {name}. Got type: {type(mm_input)}") - if isinstance(mm_input, torch.Tensor): - if dim == 0: - return mm_input.reshape(-1, *mm_input.shape[2:]) - return torch.concat(list(mm_input), dim=dim) - else: - return torch.concat(mm_input, dim=dim) - +class Qwen2_5OmniConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixinBase): def _parse_and_validate_audio_input(self, **kwargs: object) -> Qwen2_5OmniAudioFeatureInputs | None: input_audio_features = kwargs.pop("input_audio_features", None) audio_feature_lengths = kwargs.pop("audio_feature_lengths", None) feature_attention_mask = kwargs.pop("feature_attention_mask", None) if input_audio_features is None: return None - input_audio_features = self._validate_and_reshape_mm_tensor(input_audio_features, "input_audio_features", dim=1) - if feature_attention_mask is not None: - feature_attention_mask = self._validate_and_reshape_mm_tensor( - feature_attention_mask, "feature_attention_mask" - ) - if not isinstance(input_audio_features, (torch.Tensor, list)): - raise ValueError(f"Incorrect type of audio input features. Got type: {type(input_audio_features)}") + if ( + input_audio_features is not None + and isinstance(input_audio_features, torch.Tensor) + and input_audio_features.ndim == 3 + ): + input_audio_features = input_audio_features.reshape(-1, input_audio_features.shape[-1]) + elif input_audio_features is not None and isinstance(input_audio_features, list): + input_audio_features = torch.cat(input_audio_features, dim=-1) + if ( + audio_feature_lengths is not None + and isinstance(audio_feature_lengths, torch.Tensor) + and audio_feature_lengths.ndim == 2 + ): + audio_feature_lengths = audio_feature_lengths.reshape(-1) + elif audio_feature_lengths is not None and isinstance(audio_feature_lengths, list): + audio_feature_lengths = torch.cat(audio_feature_lengths, dim=-1) + if ( + feature_attention_mask is not None + and isinstance(feature_attention_mask, torch.Tensor) + and feature_attention_mask.ndim == 3 + ): + feature_attention_mask = feature_attention_mask.reshape(-1, feature_attention_mask.shape[-1]) + elif feature_attention_mask is not None and isinstance(feature_attention_mask, list): + feature_attention_mask = torch.cat(feature_attention_mask, dim=-1) return Qwen2_5OmniAudioFeatureInputs( type="audio_features", input_features=input_audio_features, @@ -706,26 +108,24 @@ class Qwen2_5OmniConditionalGenerationMixin: if pixel_values is None and image_embeds is None: return None - + if pixel_values is not None and isinstance(pixel_values, torch.Tensor) and pixel_values.ndim == 3: + pixel_values = pixel_values.reshape(-1, pixel_values.shape[-1]) + if image_embeds is not None and isinstance(image_embeds, torch.Tensor) and image_embeds.ndim == 3: + image_embeds = image_embeds.reshape(-1, image_embeds.shape[-1]) + if image_grid_thw is not None and isinstance(image_grid_thw, torch.Tensor) and image_grid_thw.ndim == 3: + image_grid_thw = image_grid_thw.reshape(-1, image_grid_thw.shape[-1]) if pixel_values is not None: - pixel_values = self._validate_and_reshape_mm_tensor(pixel_values, "image pixel values") - image_grid_thw = self._validate_and_reshape_mm_tensor(image_grid_thw, "image grid_thw") - - if not isinstance(pixel_values, (torch.Tensor, list)): - raise ValueError(f"Incorrect type of image pixel values. Got type: {type(pixel_values)}") - return Qwen2_5_VLImagePixelInputs( - type="pixel_values", pixel_values=pixel_values, image_grid_thw=image_grid_thw + type="pixel_values", + pixel_values=pixel_values, + image_grid_thw=image_grid_thw, ) if image_embeds is not None: - image_embeds = self._validate_and_reshape_mm_tensor(image_embeds, "image embeds") - image_grid_thw = self._validate_and_reshape_mm_tensor(image_grid_thw, "image grid_thw") - - if not isinstance(image_embeds, torch.Tensor): - raise ValueError(f"Incorrect type of image embeddings. Got type: {type(image_embeds)}") return Qwen2_5_VLImageEmbeddingInputs( - type="image_embeds", image_embeds=image_embeds, image_grid_thw=image_grid_thw + type="image_embeds", + image_embeds=image_embeds, + image_grid_thw=image_grid_thw, ) def _parse_and_validate_video_input( @@ -739,10 +139,17 @@ class Qwen2_5OmniConditionalGenerationMixin: if pixel_values_videos is None and video_embeds is None: return None + if ( + pixel_values_videos is not None + and isinstance(pixel_values_videos, torch.Tensor) + and pixel_values_videos.ndim == 3 + ): + pixel_values_videos = pixel_values_videos.reshape(-1, pixel_values_videos.shape[-1]) + if video_grid_thw is not None and isinstance(video_grid_thw, torch.Tensor) and video_grid_thw.ndim == 3: + video_grid_thw = video_grid_thw.reshape(-1, video_grid_thw.shape[-1]) + if video_embeds is not None and isinstance(video_embeds, torch.Tensor) and video_embeds.ndim == 3: + video_embeds = video_embeds.reshape(-1, video_embeds.shape[-1]) if pixel_values_videos is not None: - pixel_values_videos = self._validate_and_reshape_mm_tensor(pixel_values_videos, "video pixel values") - video_grid_thw = self._validate_and_reshape_mm_tensor(video_grid_thw, "video grid_thw") - return Qwen2_5_VLVideoPixelInputs( type="pixel_values_videos", pixel_values_videos=pixel_values_videos, @@ -750,79 +157,14 @@ 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") - if not isinstance(video_embeds, torch.Tensor): raise ValueError(f"Incorrect type of video embeddings. Got type: {type(video_embeds)}") return Qwen2_5_VLVideoEmbeddingInputs( - type="video_embeds", video_embeds=video_embeds, video_grid_thw=video_grid_thw + type="video_embeds", + video_embeds=video_embeds, + video_grid_thw=video_grid_thw, ) - def _process_audio_input( - self, - audio_input: Qwen2_5OmniAudioFeatureInputs, - audio_hashes: list[str] = None, - cached_audio_features: torch.Tensor = None, - ) -> torch.Tensor: - input_features = audio_input["input_features"] - audio_feature_lengths = audio_input["audio_feature_lengths"] - if input_features.ndim == 3: - assert input_features.shape[0] == 1 - input_features = input_features.squeeze(0) - if audio_feature_lengths.ndim == 2: - assert audio_feature_lengths.shape[0] == 1 or audio_feature_lengths.shape[1] == 1 - if audio_feature_lengths.shape[0] == 1: - audio_feature_lengths = audio_feature_lengths.squeeze(0) - else: - audio_feature_lengths = audio_feature_lengths.squeeze(1) - - audio_feat_lengths, audio_output_lengths = self.audio_tower._get_feat_extract_output_lengths( - audio_feature_lengths - ) - - audio_outputs = self.audio_tower( - input_features.to(self.audio_tower.dtype), - feature_lens=audio_feature_lengths, - aftercnn_lens=audio_feat_lengths, - ) - return audio_outputs.last_hidden_state.split(audio_output_lengths.tolist()) - - def _process_image_input(self, image_input: Qwen2_5_VLImageInputs) -> tuple[torch.Tensor, ...]: - if image_input["type"] == "image_embeds": - return image_input["image_embeds"].type(self.visual.dtype) - - grid_thw = image_input["image_grid_thw"] - assert grid_thw.ndim == 2 - - pixel_values = image_input["pixel_values"].type(self.visual.dtype) - image_embeds = self.visual(pixel_values, grid_thw=grid_thw) - # Split concatenated embeddings for each image item. - merge_size = self.visual.spatial_merge_size - sizes = grid_thw.prod(-1) // merge_size // merge_size - - return image_embeds.split(sizes.tolist()) - - def _process_video_input( - self, - video_input: Qwen2_5_VLVideoInputs, - video_hashes: list[str] = None, - cached_video_embeds: torch.Tensor = None, - ) -> torch.Tensor: - if video_input["type"] == "video_embeds": - return video_input["video_embeds"].type(self.visual.dtype) - - grid_thw = video_input["video_grid_thw"] - assert grid_thw.ndim == 2 - - pixel_values_videos = video_input["pixel_values_videos"].type(self.visual.dtype) - video_embeds = self.visual(pixel_values_videos, grid_thw=grid_thw) - # Split concatenated embeddings for each video item. - merge_size = self.visual.spatial_merge_size - sizes = grid_thw.prod(-1) // merge_size // merge_size - - return video_embeds.split(sizes.tolist()) - @MULTIMODAL_REGISTRY.register_processor( Qwen2_5OmniThinkerMultiModalProcessor, @@ -830,8 +172,15 @@ class Qwen2_5OmniConditionalGenerationMixin: dummy_inputs=Qwen2_5OmniThinkerDummyInputsBuilder, ) class Qwen2_5OmniThinkerForConditionalGeneration( - nn.Module, SupportsMultiModal, SupportsPP, Qwen2_5OmniConditionalGenerationMixin, SupportsMRoPE + nn.Module, + SupportsMultiModal, + SupportsPP, + SupportsLoRA, + SupportsMRoPE, + Qwen2_5OmniConditionalGenerationMixin, ): + merge_by_field_config = True + hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "thinker.lm_head.": "language_model.lm_head.", @@ -839,6 +188,22 @@ class Qwen2_5OmniThinkerForConditionalGeneration( "thinker.": "", } ) + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "attn.qkv": [ + "attn.q", + "attn.k", + "attn.v", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: @@ -853,6 +218,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() + self.vllm_config = vllm_config thinker_config: Qwen2_5OmniThinkerConfig = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config @@ -868,7 +234,8 @@ 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." ) if multimodal_config.get_limit_per_prompt("audio"): @@ -882,10 +249,10 @@ class Qwen2_5OmniThinkerForConditionalGeneration( norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), quant_config=quant_config, prefix=maybe_prefix(prefix, "visual"), - # attn_backend_override=None, ) else: self.visual = None + self.quant_config = quant_config self.language_model = init_vllm_registered_model( vllm_config=vllm_config, @@ -913,13 +280,175 @@ class Qwen2_5OmniThinkerForConditionalGeneration( def get_language_model(self) -> torch.nn.Module: return self.language_model - def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings: + def get_mrope_input_positions( + self, + input_tokens: list[int], + mm_features: list[MultiModalFeatureSpec], + ) -> tuple[torch.Tensor, int]: + """ + Example: + + (V_i are vision position ids, A_i are audio position ids) + + |V_1 ... V_n|A_1 ... A_n|V_n+1 ... V_2n|A_n+1 ... A_2n|... + |vision chunk 1|audio chunk 1|vision chunk 2|audio chunk 2 |... + """ + kwargs = MultiModalFeatureSpec.gather_kwargs( + mm_features, + { + "image_grid_thw", + "video_grid_thw", + "second_per_grid_ts", + "audio_feature_lengths", + "use_audio_in_video", + }, + ) + image_grid_thw = kwargs.get("image_grid_thw", []) + video_grid_thw = kwargs.get("video_grid_thw", []) + second_per_grid_ts = kwargs.get("second_per_grid_ts", []) + audio_feature_lengths = kwargs.get("audio_feature_lengths", []) + use_audio_in_video = any(kwargs.get("use_audio_in_video", [])) + + image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(image_grid_thw) + video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)(video_grid_thw) + + # TODO(fyabc): refactor and share more code with + # _vl_get_input_positions_tensor. + + thinker_config = self.config + audio_token_id = thinker_config.audio_token_index + image_token_id = thinker_config.image_token_index + video_token_id = thinker_config.video_token_index + audio_start_token_id = thinker_config.audio_start_token_id + audio_end_token_id = thinker_config.audio_end_token_id + vision_start_token_id = thinker_config.vision_start_token_id + vision_end_token_id = thinker_config.vision_end_token_id + seconds_per_chunk = thinker_config.seconds_per_chunk + spatial_merge_size = thinker_config.vision_config.spatial_merge_size + tokens_per_second = getattr(thinker_config.vision_config, "tokens_per_second", 25) + + src_item = input_tokens + audio_seqlens = audio_feature_lengths + if not second_per_grid_ts: + second_per_grid_ts = [1] * video_grid_thw.shape[0] + audio_idx = 0 + video_idx = 0 + image_idx = 0 + new_src_item: list[int] = [] + llm_pos_ids_list: list[torch.Tensor] = [] + + idx = 0 + while idx < len(src_item): + new_src_item_len = len(new_src_item) + start_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0 + if src_item[idx] not in [audio_token_id, video_token_id, image_token_id]: + if use_audio_in_video and idx > 0: + if src_item[idx] == vision_end_token_id and src_item[idx - 1] == audio_end_token_id: + # processing the <|audio_eos|> before <|vision_eos|> + start_idx -= 1 + elif src_item[idx] == audio_start_token_id and src_item[idx - 1] == vision_start_token_id: + # processing the <|audio_bos|> after <|vision_eos|> + start_idx -= 1 + new_src_item.append(src_item[idx]) + llm_pos_ids = torch.tensor([start_idx], dtype=torch.long).expand(3, -1) + llm_pos_ids_list.append(llm_pos_ids) + elif src_item[idx] == audio_token_id: + assert audio_seqlens is not None + audio_seqlen = audio_seqlens[audio_idx] + place_num = ((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1 + new_src_item.extend([audio_token_id] * place_num) + llm_pos_ids = torch.arange(place_num).expand(3, -1) + start_idx + llm_pos_ids_list.append(llm_pos_ids) + audio_idx += 1 + elif src_item[idx] == image_token_id: + grid_t = image_grid_thw[image_idx][0] + grid_hs = image_grid_thw[:, 1] + grid_ws = image_grid_thw[:, 2] + t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long() + llm_pos_ids = get_llm_pos_ids_for_vision( + start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws + ) + llm_pos_ids_list.append(llm_pos_ids) + vision_seqlen = image_grid_thw[image_idx].prod() // (spatial_merge_size**2) + new_src_item.extend([image_token_id] * vision_seqlen) + image_idx += 1 + elif src_item[idx] == video_token_id and not use_audio_in_video: + grid_t = video_grid_thw[video_idx][0] + grid_hs = video_grid_thw[:, 1] + grid_ws = video_grid_thw[:, 2] + t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long() + llm_pos_ids = get_llm_pos_ids_for_vision( + start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws + ) + llm_pos_ids_list.append(llm_pos_ids) + vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2) + new_src_item.extend([video_token_id] * vision_seqlen) + video_idx += 1 + else: + # read audio from video + assert audio_seqlens is not None + audio_seqlen = audio_seqlens[audio_idx] + vision_seqlen = video_grid_thw[video_idx].prod() // (spatial_merge_size**2) + grid_t = video_grid_thw[video_idx][0] + grid_h = video_grid_thw[video_idx][1] + grid_w = video_grid_thw[video_idx][2] + grid_hs = video_grid_thw[:, 1] + grid_ws = video_grid_thw[:, 2] + t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk) + t_index = (torch.arange(grid_t) * second_per_grid_ts[video_idx] * tokens_per_second).long() + t_index_split_chunk = split_list_into_ranges(t_index, t_ntoken_per_chunk) + place_num = (((audio_seqlen - 1) // 2 + 1 - 2) // 2 + 1) + 2 + pure_audio_len = place_num - 2 + added_audio_len = 0 + audio_llm_pos_ids_list: list[torch.Tensor] = [] + for t_chunk in t_index_split_chunk: + vision_ntoken_per_chunk = len(t_chunk) * grid_h * grid_w // (spatial_merge_size**2) + new_src_item.extend([video_token_id] * vision_ntoken_per_chunk) + vision_llm_pos_ids_list = get_llm_pos_ids_for_vision( + start_idx, + video_idx, + spatial_merge_size, + t_chunk, + grid_hs, + grid_ws, + ).split(1, dim=1) + llm_pos_ids_list.extend(vision_llm_pos_ids_list) + new_src_item.extend(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) * [audio_token_id]) + audio_start_idx = ( + start_idx if len(audio_llm_pos_ids_list) == 0 else audio_llm_pos_ids_list[-1][0].item() + 1 + ) + if min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) > 0: + audio_llm_pos_ids_list = ( + torch.arange(min(t_ntoken_per_chunk, pure_audio_len - added_audio_len)).expand(3, -1) + + audio_start_idx + ).split(1, dim=1) + else: + audio_llm_pos_ids_list = [] + added_audio_len += min(t_ntoken_per_chunk, pure_audio_len - added_audio_len) + llm_pos_ids_list.extend(audio_llm_pos_ids_list) + if added_audio_len < pure_audio_len: + new_src_item.extend((pure_audio_len - added_audio_len) * [audio_token_id]) + audio_llm_pos_ids_list = ( + torch.arange(pure_audio_len - added_audio_len).expand(3, -1) + llm_pos_ids_list[-1].max() + 1 + ).split(1, dim=1) + llm_pos_ids_list.extend(audio_llm_pos_ids_list) + audio_idx += 1 + video_idx += 1 + # move to the next token + idx += len(new_src_item) - new_src_item_len + + llm_positions = torch.cat(llm_pos_ids_list, dim=1) + mrope_position_delta = torch.cat(llm_pos_ids_list, dim=1).max() + 1 - len(src_item) + + return llm_positions, mrope_position_delta + + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings: mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) if not mm_input_by_modality: return [] # The result multimodal_embeddings is tuple of tensors, with each - # tensor correspoending to a multimodal data item (image or video). + # tensor corresponding to a multimodal data item (image or video). multimodal_embeddings: tuple[torch.Tensor, ...] = () # NOTE: It is important to iterate over the keys in this dictionary @@ -927,76 +456,36 @@ class Qwen2_5OmniThinkerForConditionalGeneration( for modality in mm_input_by_modality: multimodal_input = mm_input_by_modality[modality] if modality == "image": - vision_embeddings = self._process_image_input(multimodal_input) - multimodal_embeddings += vision_embeddings + image_embeddings = self._process_image_input(multimodal_input) + multimodal_embeddings += tuple(image_embeddings) if modality == "video": video_embeddings = self._process_video_input(multimodal_input) - multimodal_embeddings += video_embeddings + multimodal_embeddings += tuple(video_embeddings) if modality == "audio": audio_embeddings = self._process_audio_input(multimodal_input) - multimodal_embeddings += audio_embeddings + multimodal_embeddings += tuple(audio_embeddings) return multimodal_embeddings - def get_input_embeddings( + # TODO (ywang96): support overlapping modality embeddings so that + # `use_audio_in_video` will work on V1. + def embed_input_ids( self, input_ids: torch.Tensor, multimodal_embeddings: MultiModalEmbeddings | None = None, + *, + is_multimodal: torch.Tensor | None = None, + handle_oov_mm_token: bool = False, ) -> torch.Tensor: - inputs_embeds = self.language_model.get_input_embeddings(input_ids) - if multimodal_embeddings is not None and len(multimodal_embeddings) != 0: - # TODO (ywang96): support overlapping modalitiy embeddings so that - # `use_audio_in_video` will work on V1. - inputs_embeds = merge_multimodal_embeddings( - input_ids, - inputs_embeds, - multimodal_embeddings, - [ - self.config.image_token_index, - self.config.video_token_index, - self.config.audio_token_index, - ], - ) - return inputs_embeds + # This is to satisfy the type checker for each overload + if multimodal_embeddings is None or is_multimodal is None: + return super().embed_input_ids(input_ids) - def get_multimodal_embeddings_v0(self, **kwargs: object) -> NestedTensors | None: - audio_input = self._parse_and_validate_audio_input(**kwargs) - image_input = self._parse_and_validate_image_input(**kwargs) - video_input = self._parse_and_validate_video_input(**kwargs) - - if audio_input is None and image_input is None and video_input is None: - return None - - multimodal_embeddings: list[tuple[NestedTensors, str]] = [] - - if audio_input is not None: - audio_embeds = self._process_audio_input(audio_input) - multimodal_embeddings.append((audio_embeds, "audio")) - if image_input is not None: - image_embeds = self._process_image_input(image_input) - multimodal_embeddings.append((image_embeds, "image")) - if video_input is not None: - video_embeds = self._process_video_input(video_input) - multimodal_embeddings.append((video_embeds, "video")) - return multimodal_embeddings - - def get_input_embeddings_v0( - self, - input_ids: torch.Tensor, - multimodal_embeddings: NestedTensors | None = None, - ) -> torch.Tensor: - inputs_embeds = self.language_model.get_input_embeddings(input_ids) - if multimodal_embeddings is None or len(multimodal_embeddings) == 0: - return inputs_embeds - - for embeddings, modality in multimodal_embeddings: - if modality == "audio": - placeholder_token_id = self.config.audio_token_index - if modality == "image": - placeholder_token_id = self.config.image_token_index - if modality == "video": - placeholder_token_id = self.config.video_token_index - inputs_embeds = merge_multimodal_embeddings(input_ids, inputs_embeds, embeddings, placeholder_token_id) - return inputs_embeds + return super().embed_input_ids( + input_ids, + multimodal_embeddings=multimodal_embeddings, + is_multimodal=is_multimodal, + handle_oov_mm_token=handle_oov_mm_token, + ) def forward( self, @@ -1009,51 +498,38 @@ class Qwen2_5OmniThinkerForConditionalGeneration( if intermediate_tensors is not None: inputs_embeds = None - # NOTE: In v1, inputs_embeds is always generated at model runner, this - # condition is for v0 compatibility. - elif inputs_embeds is None: - multimodal_embeddings = self.get_multimodal_embeddings(**kwargs) - inputs_embeds = self.get_input_embeddings(input_ids, multimodal_embeddings) - text_inputs_embeds = self.get_input_embeddings( - input_ids, - ( - [(torch.zeros_like(embeddings), modality) for embeddings, modality in multimodal_embeddings] - if multimodal_embeddings is not None - else None - ), - ) - input_ids = None - else: - text_inputs_embeds = inputs_embeds - hidden_states = self.language_model.model( input_ids, positions, intermediate_tensors, inputs_embeds=inputs_embeds ) - return text_inputs_embeds, hidden_states.unsqueeze(0) # (1, S, D) + return hidden_states - def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None: + def compute_logits( + self, + hidden_states: torch.Tensor, + ) -> torch.Tensor | None: return self.language_model.compute_logits(hidden_states) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + skip_prefixes = ["talker.", "token2wav."] + if self.audio_tower is None: + skip_prefixes.extend(["audio_tower."]) + if self.visual is None: + skip_prefixes.extend(["visual."]) + loader = AutoWeightsLoader( self, - skip_prefixes=["talker.", "token2wav."], + skip_prefixes=skip_prefixes, ) loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - # Log load summary - try: - total_bytes = 0 - for name, param in self.named_parameters(): - if param is not None and param.data is not None: - total_bytes += param.data.numel() * param.data.element_size() - device = next(self.parameters()).device - logger.info( - "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s", - self.__class__.__name__, - True, - total_bytes / (1024**2), - str(device), - ) - except Exception: - pass + return loaded_weights + + def get_mm_mapping(self) -> MultiModelKeys: + """ + Get the module prefix in multimodal models + """ + return MultiModelKeys.from_string_field( + language_model="language_model", + connector="merger.", + tower_model=["visual.", "audio_tower."], + ) diff --git a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py index f6f05357a..ce9819b3e 100644 --- a/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py +++ b/vllm_omni/model_executor/models/qwen2_5_omni/qwen2_old.py @@ -3,7 +3,10 @@ from collections.abc import Iterable import torch from torch import nn from transformers import Qwen2Config -from vllm.attention import Attention, AttentionType +from vllm.attention.backends.abstract import ( + AttentionType, +) +from vllm.attention.layer import Attention from vllm.compilation.decorators import support_torch_compile from vllm.config import CacheConfig, VllmConfig from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size @@ -125,12 +128,15 @@ class Qwen2Attention(nn.Module): prefix=f"{prefix}.o_proj", ) - self.rotary_emb = get_rope( - self.head_dim, + self.rotary_pos_emb = get_rope( + head_size=self.head_dim, rotary_dim=self.head_dim, max_position=max_position, - base=self.rope_theta, - rope_scaling=rope_scaling, + is_neox_style=True, + rope_parameters={ + "base": self.rope_theta, + **rope_scaling, + }, ) self.attn = Attention( self.num_heads, @@ -150,7 +156,7 @@ class Qwen2Attention(nn.Module): ) -> torch.Tensor: qkv, _ = self.qkv_proj(hidden_states) q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1) - q, k = self.rotary_emb(positions, q, k) + q, k = self.rotary_pos_emb(positions, q, k) attn_output = self.attn(q, k, v) output, _ = self.o_proj(attn_output) return output @@ -294,7 +300,7 @@ class Qwen2Model(nn.Module): else: self.norm = PPMissingLayer() - def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: return self.embed_tokens(input_ids) def forward( @@ -308,7 +314,7 @@ class Qwen2Model(nn.Module): if inputs_embeds is not None: hidden_states = inputs_embeds else: - hidden_states = self.get_input_embeddings(input_ids) + hidden_states = self.embed_input_ids(input_ids) residual = None else: assert intermediate_tensors is not None @@ -420,8 +426,8 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): self.make_empty_intermediate_tensors = self.model.make_empty_intermediate_tensors - def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor: - return self.model.get_input_embeddings(input_ids) + def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor: + return self.model.embed_input_ids(input_ids) def forward( self, diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py index 12faccd2f..1a6cbae98 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -18,6 +18,12 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsMultiModal, SupportsPP +from vllm.model_executor.models.qwen3_omni_moe_thinker import ( + Qwen3OmniMoeConditionalGenerationMixin, + Qwen3OmniMoeThinkerDummyInputsBuilder, + Qwen3OmniMoeThinkerMultiModalProcessor, + Qwen3OmniMoeThinkerProcessingInfo, +) from vllm.model_executor.models.utils import init_vllm_registered_model, maybe_prefix from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors @@ -28,13 +34,6 @@ from vllm.v1.sample.sampler import Sampler from vllm_omni.model_executor.models.output_templates import OmniOutput from vllm_omni.model_executor.models.utils import add_prefix_to_loaded_weights -from .qwen3_omni_moe_thinker import ( - Qwen3OmniMoeConditionalGenerationMixin, - Qwen3OmniMoeThinkerDummyInputsBuilder, - Qwen3OmniMoeThinkerMultiModalProcessor, - Qwen3OmniMoeThinkerProcessingInfo, -) - # Special token IDs for Qwen3 Omni MoE # Reference: https://huggingface.co/Qwen/Qwen3-Omni-30B-A3B-Instruct/blob/main/tokenizer_config.json @@ -226,20 +225,21 @@ class Qwen3OmniMoeForConditionalGeneration( return self.model.sampler return Sampler() - def get_input_embeddings( + def embed_input_ids( self, input_ids: torch.Tensor, multimodal_embeddings=None, + is_multimodal=None, ) -> torch.Tensor: - """Get input embeddings for the active model stage.""" if self.model_stage == "code2wav": - # Code2wav doesn't use text embeddings return torch.zeros_like(input_ids).reshape(-1, 1).repeat(1, self.vllm_config.model_config.get_hidden_size()) - return self.model.get_input_embeddings(input_ids, multimodal_embeddings) + return self.model.embed_input_ids( + input_ids=input_ids, multimodal_embeddings=multimodal_embeddings, is_multimodal=is_multimodal + ) - def get_multimodal_embeddings(self, **kwargs): + def embed_multimodal(self, **kwargs): """Delegate to active model for multimodal processing.""" - return self.model.get_multimodal_embeddings(**kwargs) + return self.model.embed_multimodal(**kwargs) # ==================== Forward Pass ==================== def _get_talker_suppressed_tokens(self): @@ -281,16 +281,16 @@ class Qwen3OmniMoeForConditionalGeneration( # ========== Stage 1: Thinker ========== if self.model_stage == "thinker": # Normalize to batched inputs if needed - added_batch_dim = False + _added_batch_dim = False if input_ids is not None and input_ids.ndim == 1: input_ids = input_ids.unsqueeze(0) - added_batch_dim = True + _added_batch_dim = True if positions is not None and positions.ndim == 1: positions = positions.unsqueeze(0) - added_batch_dim = True + _added_batch_dim = True if inputs_embeds is not None and inputs_embeds.ndim == 2: inputs_embeds = inputs_embeds.unsqueeze(0) - added_batch_dim = True + _added_batch_dim = True thinker_dev = self._module_device(self.thinker) @@ -301,7 +301,7 @@ class Qwen3OmniMoeForConditionalGeneration( dtype=torch.long, device=thinker_dev, ).unsqueeze(0) - added_batch_dim = True + _added_batch_dim = True # Move to thinker device if input_ids is not None and input_ids.device != thinker_dev: @@ -337,7 +337,7 @@ class Qwen3OmniMoeForConditionalGeneration( device=self._module_device(self.thinker), dtype=torch.long, ) - thinker_tts_embeds = self.thinker.get_input_embeddings(tts_tokens) # [1,3,thinker_hidden] + thinker_tts_embeds = self.thinker.embed_input_ids(tts_tokens) # [1,3,thinker_hidden] if ( isinstance(thinker_tts_embeds, torch.Tensor) and thinker_tts_embeds.ndim == 3 @@ -353,7 +353,7 @@ class Qwen3OmniMoeForConditionalGeneration( # Return text-only output (with multimodal sidecar) return OmniOutput( - text_hidden_states=(text_hidden_states.squeeze(0) if added_batch_dim else text_hidden_states), + text_hidden_states=(text_hidden_states.reshape(-1, text_hidden_states.shape[-1])), multimodal_outputs=multimodal_outputs, ) @@ -379,7 +379,7 @@ class Qwen3OmniMoeForConditionalGeneration( # Ensure we have base embeddings when only ids are provided if inputs_embeds is None and input_ids is not None: - inputs_embeds = self.talker.get_input_embeddings(input_ids) + inputs_embeds = self.talker.embed_input_ids(input_ids) # ------- Request-scoped additional information (no cross-request concat) ------- request_ids: list[str] | None = kwargs.get("request_ids") # ordered @@ -748,8 +748,6 @@ class Qwen3OmniMoeForConditionalGeneration( - hidden_projection: Used at runtime for multimodal hidden states (audio/image/video) from thinker's last layer, not needed for special token initialization """ - # Get embeddings from both models - # self.thinker_embedding = self.thinker.model.get_input_embeddings() self.talker_embedding = self._load_talker_embedding() # Get configuration @@ -960,9 +958,7 @@ class Qwen3OmniMoeForConditionalGeneration( device=tts_pad_embed.device, dtype=torch.bfloat16, ), - self.talker.get_input_embeddings(codec_special_tokens).to( - device=tts_pad_embed.device, dtype=torch.bfloat16 - ), + self.talker.embed_input_ids(codec_special_tokens).to(device=tts_pad_embed.device, dtype=torch.bfloat16), ), dim=0, ) @@ -1059,7 +1055,7 @@ class Qwen3OmniMoeForConditionalGeneration( logitsprocs=LogitsProcessors(), ) # Use active model for logits computation - logits = self.model.compute_logits(hidden_states, sampling_metadata) # V, d + logits = self.model.compute_logits(hidden_states) # V, d # Talker: suppress tokens by setting their probability to ~1e-9 (finite very small), # implemented by assigning their logits to log(1e-9). diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py index 91231d0fb..48fde0564 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_talker.py @@ -23,6 +23,7 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import ( Qwen2_5OmniThinkerDummyInputsBuilder, ) from vllm.model_executor.models.qwen3_moe import Qwen3MoeMLP +from vllm.model_executor.models.qwen3_omni_moe_thinker import Qwen3Omni_VisionTransformer from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, @@ -31,14 +32,12 @@ from vllm.model_executor.models.utils import ( ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.sequence import IntermediateTensors -from vllm.v1.sample.metadata import SamplingMetadata from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_code_predictor_mtp import ( Qwen3OmniMoeTalkerCodePredictor, ) from vllm_omni.model_executor.models.qwen3_omni.qwen3_omni_moe_thinker import ( Qwen3MoeLLMForCausalLM, - Qwen3Omni_VisionTransformer, Qwen3OmniMoeConditionalGenerationMixin, Qwen3OmniMoeThinkerMultiModalProcessor, Qwen3OmniMoeThinkerProcessingInfo, @@ -106,9 +105,11 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() talker_config: Qwen3OmniMoeTalkerConfig = vllm_config.model_config.hf_config + talker_config.text_config.rope_parameters = talker_config.text_config.rope_scaling + talker_config.text_config.rope_parameters["rope_theta"] = talker_config.text_config.rope_theta self.quant_config = vllm_config.quant_config self.prefix = prefix - + self.vllm_config = vllm_config self.config = talker_config self.vocab_size = talker_config.text_config.vocab_size self.router_aux_loss_coef = talker_config.text_config.router_aux_loss_coef @@ -193,7 +194,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( pos_codes = [layer0_code] # Start with layer 0: [batch, 1] # Initial input: [last_talker_hidden, layer0_embed] - layer0_embed = self.get_input_embeddings(layer0_code) + layer0_embed = self.embed_input_ids(layer0_code) prev_embed = layer0_embed # Track previous layer embedding try: current_input = torch.cat([last_talker_hidden, prev_embed], dim=1) # [batch, 2, hidden_size] @@ -260,7 +261,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( for pos in range(seq_len): # Layer 0 embedding layer0_code = result_codes[:, 0, pos : pos + 1] # [batch, 1] - layer0_embed = self.get_input_embeddings(layer0_code) # [batch, 1, hidden_size] + layer0_embed = self.embed_input_ids(layer0_code) # [batch, 1, hidden_size] # mid layers hidden states (from CodePredictor) mid_residual_hiddens = middle_hidden_states[pos] # [batch, num_code_groups-2, hidden_size] @@ -440,7 +441,6 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( def compute_logits( self, hidden_states: torch.Tensor, - sampling_metadata: SamplingMetadata = None, ) -> torch.Tensor | None: """Compute logits for audio codec codes (layer 0 of RVQ). @@ -471,7 +471,7 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(**kwargs) return mm_input_by_modality - def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings | None: + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) if not mm_input_by_modality: return [] @@ -504,9 +504,14 @@ class Qwen3OmniMoeTalkerForConditionalGeneration( multimodal_embeddings += tuple(audio_embeddings) return multimodal_embeddings - def get_input_embeddings(self, input_ids: torch.Tensor, multimodal_embeddings: MultiModalEmbeddings | None = None): + def embed_input_ids( + self, + input_ids: torch.Tensor, + multimodal_embeddings: MultiModalEmbeddings | None = None, + is_multimodal: bool = False, + ): """Get the input embedding layer (for codec tokens).""" - return self.language_model.get_input_embeddings(input_ids, multimodal_embeddings) + return self.language_model.embed_input_ids(input_ids) def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: """Load weights for the talker model. @@ -695,10 +700,9 @@ class Qwen3OmniMoeModel(Qwen3MoeLLMForCausalLM): return forward_with_shared_expert - def get_input_embeddings( + def embed_input_ids( self, input_ids: torch.Tensor, - multimodal_embeddings: MultiModalEmbeddings | None = None, - generation_steps=None, + **kwargs: object, ) -> torch.Tensor: return self.model.codec_embedding(input_ids) diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py index b13cf6aa5..9a98712e0 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni_moe_thinker.py @@ -22,15 +22,13 @@ # limitations under the License. """Inference-only Qwen3-Omni-Moe model (thinker part).""" -import math -from collections.abc import Callable, Iterable, Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from functools import partial from typing import Any import numpy as np import torch import torch.nn as nn -import torch.nn.functional as F from packaging.version import Version from transformers import PretrainedConfig from transformers import __version__ as TRANSFORMERS_VERSION @@ -46,52 +44,44 @@ from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import ( Qwen3OmniMoeProcessor, ) from transformers.models.whisper import WhisperFeatureExtractor -from vllm.attention.utils.fa_utils import get_flash_attn_version from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.distributed import get_pp_group from vllm.logger import init_logger -from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY -from vllm.model_executor.layers.linear import ( - ColumnParallelLinear, - ReplicatedLinear, - RowParallelLinear, -) from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead -from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, + SupportsMRoPE, SupportsMultiModal, SupportsPP, ) from vllm.model_executor.models.qwen2_5_omni_thinker import ( - Qwen2_5OmniConditionalGenerationMixin, + Qwen2_5OmniAudioFeatureInputs, Qwen2_5OmniThinkerDummyInputsBuilder, Qwen2_5OmniThinkerMultiModalProcessor, ) from vllm.model_executor.models.qwen2_5_vl import ( - Qwen2_5_VisionAttention, - Qwen2_5_VisionRotaryEmbedding, Qwen2_5_VLProcessingInfo, ) -from vllm.model_executor.models.qwen2_audio import Qwen2AudioInputs, Qwen2AudioProcessingInfo +from vllm.model_executor.models.qwen2_audio import Qwen2AudioProcessingInfo from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM +from vllm.model_executor.models.qwen3_moe import Qwen3MoeModel as _Qwen3MoeLLMModel +from vllm.model_executor.models.qwen3_omni_moe_thinker import ( + Qwen3Omni_VisionTransformer, + _get_feat_extract_output_lengths, +) from vllm.model_executor.models.utils import ( AutoWeightsLoader, WeightsMapper, _merge_multimodal_embeddings, - flatten_bn, maybe_prefix, ) from vllm.model_executor.models.vision import ( - get_vit_attn_backend, + get_llm_pos_ids_for_vision, ) from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import ( - MultiModalKwargsItems, -) +from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItems from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems from vllm.multimodal.processing import ( MultiModalPromptUpdates, @@ -100,504 +90,20 @@ from vllm.multimodal.processing import ( PromptUpdate, PromptUpdateDetails, ) -from vllm.platforms import _Backend from vllm.sequence import IntermediateTensors -from vllm.v1.sample.metadata import SamplingMetadata -from .qwen3_moe import Qwen3MoeModel +from vllm_omni.model_executor.models.qwen2_5_omni.qwen2_5_omni_thinker import ( + Qwen2_5OmniConditionalGenerationMixin, +) try: import flash_attn except (ImportError, ModuleNotFoundError): flash_attn = None -from torch import Tensor - logger = init_logger(__name__) -# Local helper functions (not available in vLLM) -def conv3d_to_linear_weight(weight: torch.Tensor) -> torch.Tensor: - """Convert 3D convolutional weight to linear weight for patch embedding.""" - # Weight shape: (out_channels, in_channels, t, h, w) - # Need to flatten spatial dimensions: (out_channels, in_channels * t * h * w) - flattened_weight = weight.flatten(1) # Flatten all dims except out_channels - return flattened_weight - - -def get_llm_pos_ids_for_vision( - start_idx: int, - vision_idx: int, - spatial_merge_size: int, - t_index: list[torch.Tensor], - grid_hs: list[torch.Tensor], - grid_ws: list[torch.Tensor], -) -> torch.Tensor: - """Generate 3D position IDs for vision tokens (temporal, height, width).""" - llm_grid_h = grid_hs[vision_idx] // spatial_merge_size - llm_grid_w = grid_ws[vision_idx] // spatial_merge_size - h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(len(t_index), -1, llm_grid_w).flatten().float() - w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(len(t_index), llm_grid_h, -1).flatten().float() - t_index_tensor = torch.Tensor(t_index).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten().float() - llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index]) - llm_pos_ids = llm_pos_ids + start_idx - return llm_pos_ids - - -# Type aliases for audio inputs -Qwen2_5OmniAudioFeatureInputs = Qwen2AudioInputs - - -def _get_feat_extract_output_lengths(input_lengths: torch.Tensor): - input_lengths_leave = input_lengths % 100 - feat_lengths = (input_lengths_leave - 1) // 2 + 1 - output_lengths = ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13 - return feat_lengths, output_lengths - - -class Qwen3_VisionPatchEmbed(nn.Module): - def __init__( - self, - patch_size: int = 14, - temporal_patch_size: int = 2, - in_channels: int = 3, - hidden_size: int = 1152, - ) -> None: - super().__init__() - self.patch_size = patch_size - self.temporal_patch_size = temporal_patch_size - self.hidden_size = hidden_size - - kernel_size = (temporal_patch_size, patch_size, patch_size) - self.proj = ReplicatedLinear( - in_channels * math.prod(kernel_size), - hidden_size, - bias=True, - return_bias=False, - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - L, C = x.shape - x = self.proj(x) - return x - - -class Qwen3_VisionMLP(nn.Module): - def __init__( - self, - in_features: int, - hidden_features: int, - bias: bool = False, - act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ): - super().__init__() - self.linear_fc1 = ColumnParallelLinear( - in_features, - hidden_features, - bias=bias, - quant_config=quant_config, - return_bias=False, - prefix=f"{prefix}.linear_fc1", - ) - self.linear_fc2 = RowParallelLinear( - hidden_features, - in_features, - bias=bias, - quant_config=quant_config, - return_bias=False, - prefix=f"{prefix}.linear_fc2", - ) - self.act_fn = act_fn - - def forward(self, x: torch.Tensor): - mlp_output = self.linear_fc2(self.act_fn(self.linear_fc1(x))) - return mlp_output - - -class Qwen3_VisionBlock(nn.Module): - def __init__( - self, - dim: int, - num_heads: int, - mlp_hidden_dim: int, - act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu, - norm_layer: Callable[[int], nn.Module] | None = None, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - if norm_layer is None: - norm_layer = partial(nn.LayerNorm, eps=1e-6) - self.norm1 = norm_layer(dim) - self.norm2 = norm_layer(dim) - self.attn = Qwen2_5_VisionAttention( - embed_dim=dim, - num_heads=num_heads, - projection_size=dim, - quant_config=quant_config, - prefix=f"{prefix}.attn", - ) - self.mlp = Qwen3_VisionMLP( - dim, - mlp_hidden_dim, - act_fn=act_fn, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.mlp", - ) - - def forward( - self, - x: torch.Tensor, - cu_seqlens: torch.Tensor, - rotary_pos_emb: torch.Tensor, - max_seqlen: int | None = None, # Only used for Flash Attention - seqlens: list[int] | None = None, # Only used for xFormers - ) -> torch.Tensor: - x = x + self.attn( - self.norm1(x), - cu_seqlens=cu_seqlens, - rotary_pos_emb=rotary_pos_emb, - max_seqlen=max_seqlen, - seqlens=seqlens, - ) - - x = x + self.mlp(self.norm2(x)) - return x - - -class Qwen3_VisionPatchMerger(nn.Module): - def __init__( - self, - d_model: int, - context_dim: int, - norm_layer: Callable[[int], nn.Module] | None = None, - spatial_merge_size: int = 2, - use_postshuffle_norm: bool = False, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - ) -> None: - super().__init__() - self.hidden_size = context_dim * (spatial_merge_size**2) - - self.use_postshuffle_norm = use_postshuffle_norm - if self.use_postshuffle_norm: - context_dim = self.hidden_size - - if norm_layer is None: - norm_layer = partial(nn.LayerNorm, eps=1e-6) - self.use_postshuffle_norm = use_postshuffle_norm - self.ln_q = norm_layer(self.hidden_size if use_postshuffle_norm else context_dim) - self.mlp = nn.ModuleList( - [ - ColumnParallelLinear( - self.hidden_size, - self.hidden_size, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.mlp.0", - ), - nn.GELU(), - RowParallelLinear( - self.hidden_size, - d_model, - bias=True, - quant_config=quant_config, - prefix=f"{prefix}.mlp.2", - ), - ] - ) - - def forward(self, x: torch.Tensor) -> torch.Tensor: - if self.use_postshuffle_norm: - x = self.ln_q(x.view(-1, self.hidden_size)) - else: - x = self.ln_q(x).view(-1, self.hidden_size) - - mlp_fc1, mlp_act, mlp_fc2 = self.mlp - x_parallel, _ = mlp_fc1(x) - x_parallel = mlp_act(x_parallel) - out, _ = mlp_fc2(x_parallel) - return out - - -class Qwen3Omni_VisionTransformer(nn.Module): - def __init__( - self, - vision_config, - norm_eps: float = 1e-6, - quant_config: QuantizationConfig | None = None, - prefix: str = "", - attn_backend_override: _Backend | None = None, - ) -> None: - super().__init__() - self.hidden_size = vision_config.hidden_size - self.num_heads = vision_config.num_heads - self.image_size = vision_config.image_size - self.patch_size = vision_config.patch_size - self.spatial_merge_size = vision_config.spatial_merge_size - self.spatial_merge_unit = self.spatial_merge_size**2 - self.temporal_patch_size = vision_config.temporal_patch_size - self.num_grid_per_side = self.image_size // self.patch_size - self.apply_vit_abs_pos_embed = vision_config.apply_vit_abs_pos_embed - self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes - - self.patch_embed = Qwen3_VisionPatchEmbed( - patch_size=self.patch_size, - temporal_patch_size=self.temporal_patch_size, - in_channels=vision_config.in_channels, - hidden_size=self.hidden_size, - ) - - # vit pos embedding, TODO: spatial_patch_size vs patch_size - if self.apply_vit_abs_pos_embed: - self.pos_embed = nn.Embedding(self.num_grid_per_side**2, self.hidden_size) - else: - self.pos_embed = nn.Parameter(torch.empty([1, self.num_grid_per_side**2, self.hidden_size])) - - norm_layer = partial(nn.LayerNorm, eps=norm_eps) - head_dim = self.hidden_size // self.num_heads - self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2) - - self.blocks = nn.ModuleList( - [ - Qwen3_VisionBlock( - dim=self.hidden_size, - num_heads=self.num_heads, - mlp_hidden_dim=vision_config.intermediate_size, - act_fn=_ACTIVATION_REGISTRY[vision_config.hidden_act], - norm_layer=norm_layer, - quant_config=quant_config, - prefix=f"{prefix}.blocks.{layer_idx}", - ) - for layer_idx in range(vision_config.depth) - ] - ) - self.merger = Qwen3_VisionPatchMerger( - d_model=vision_config.out_hidden_size, - context_dim=self.hidden_size, - norm_layer=norm_layer, - spatial_merge_size=self.spatial_merge_size, - quant_config=quant_config, - prefix=f"{prefix}.merger", - ) - if self.deepstack_visual_indexes is not None: - self.merger_list = nn.ModuleList( - [ - Qwen3_VisionPatchMerger( - d_model=vision_config.out_hidden_size, - context_dim=self.hidden_size, - spatial_merge_size=self.spatial_merge_size, - use_postshuffle_norm=True, - norm_layer=norm_layer, - quant_config=quant_config, - prefix=f"{prefix}.merger_list.{layer_idx}", - ) - for layer_idx in range(len(self.deepstack_visual_indexes)) - ] - ) - - self.attn_backend = get_vit_attn_backend( - head_size=head_dim, - dtype=torch.get_default_dtype(), - # attn_backend_override=attn_backend_override, - ) - if self.attn_backend != _Backend.FLASH_ATTN and get_flash_attn_version() is not None: - self.attn_backend = _Backend.FLASH_ATTN - - @property - def dtype(self) -> torch.dtype: - return self.patch_embed.proj.weight.dtype - - @property - def device(self) -> torch.device: - return self.patch_embed.proj.weight.device - - def rot_pos_emb(self, grid_thw): - pos_ids = [] - for t, h, w in grid_thw: - hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w) - hpos_ids = hpos_ids.reshape( - h // self.spatial_merge_size, - self.spatial_merge_size, - w // self.spatial_merge_size, - self.spatial_merge_size, - ) - hpos_ids = hpos_ids.permute(0, 2, 1, 3) - hpos_ids = hpos_ids.flatten() - - wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1) - wpos_ids = wpos_ids.reshape( - h // self.spatial_merge_size, - self.spatial_merge_size, - w // self.spatial_merge_size, - self.spatial_merge_size, - ) - wpos_ids = wpos_ids.permute(0, 2, 1, 3) - wpos_ids = wpos_ids.flatten() - pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1)) - pos_ids = torch.cat(pos_ids, dim=0) - max_grid_size = grid_thw[:, 1:].max() - rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size) - rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1) - return rotary_pos_emb - - def fast_pos_embed_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor: - num_grid_per_side = self.num_grid_per_side - m_size = self.spatial_merge_size - hidden_dim = self.pos_embed.embedding_dim - - outputs = [] - for t, h, w in grid_thw: - h_idxs = torch.linspace(0, num_grid_per_side - 1, h, dtype=torch.float32, device=self.device) - w_idxs = torch.linspace(0, num_grid_per_side - 1, w, dtype=torch.float32, device=self.device) - - h_floor = h_idxs.to(torch.long) - w_floor = w_idxs.to(torch.long) - h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1) - w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1) - - dh = h_idxs - h_floor - dw = w_idxs - w_floor - - # Create meshgrid view for all h, w vars - dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij") - h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij") - h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij") - h_floor_grid_idx = h_floor_grid * num_grid_per_side - h_ceil_grid_idx = h_ceil_grid * num_grid_per_side - - # original computation of weights - # w00 = (1 - dh_grid) * (1 - dw_grid) - # w01 = (1 - dh_grid) * dw_grid - # w10 = dh_grid * (1 - dw_grid) - # w11 = dh_grid * dw_grid - # we reuse w11 here to avoid duplicate - # dh_grid * dw_grid computation - w11 = dh_grid * dw_grid - w10 = dh_grid - w11 - w01 = dw_grid - w11 - w00 = 1 - dh_grid - dw_grid + w11 - - idx00 = h_floor_grid_idx + w_floor_grid - idx01 = h_floor_grid_idx + w_ceil_grid - idx10 = h_ceil_grid_idx + w_floor_grid - idx11 = h_ceil_grid_idx + w_ceil_grid - - indices = torch.stack([idx00, idx01, idx10, idx11], dim=0).reshape(4, -1) - weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1) - weights = weights.to(dtype=self.dtype, device=self.device) - - embeds = self.pos_embed(indices) - weighted_embeds = embeds * weights - p0, p1, p2, p3 = weighted_embeds.unbind(dim=0) - combined = p0 + p1 + p2 + p3 - - combined = combined.view(h * w, hidden_dim) - repeated = combined.unsqueeze(0).expand(t, -1, -1).contiguous() - repeated = repeated.view(t, h // m_size, m_size, w // m_size, m_size, hidden_dim) - repeated = repeated.permute(0, 1, 3, 2, 4, 5).reshape(-1, hidden_dim) - outputs.append(repeated) - - return torch.cat(outputs, dim=0) - - def compute_attn_mask_seqlen( - self, - cu_seqlens: torch.Tensor, - ) -> tuple[int | None, list[int] | None]: - max_seqlen, seqlens = None, None - if self.attn_backend == _Backend.FLASH_ATTN: - max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item() - elif self.attn_backend == _Backend.XFORMERS: - seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() - return max_seqlen, seqlens - - def forward( - self, - x: torch.Tensor, - grid_thw: list[list[int]], - ) -> torch.Tensor: - hidden_states = x.to(device=self.device, dtype=self.dtype) - hidden_states = self.patch_embed(hidden_states) - - if self.apply_vit_abs_pos_embed: - pos_embeds = self.fast_pos_embed_interpolate(grid_thw) - hidden_states = hidden_states + pos_embeds - rotary_pos_emb = self.rot_pos_emb(grid_thw) - - cu_seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]).cumsum( - dim=0, - dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32, - ) - cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0) - - hidden_states = hidden_states.unsqueeze(1) - rotary_pos_emb = rotary_pos_emb.to(hidden_states.device) - max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens) - - hidden_states_list = [] - deepstack_visual_indexes = self.deepstack_visual_indexes - - for layer_num, blk in enumerate(self.blocks): - hidden_states = blk( - hidden_states, - cu_seqlens=cu_seqlens, - rotary_pos_emb=rotary_pos_emb, - max_seqlen=max_seqlen, - seqlens=seqlens, - ) - if deepstack_visual_indexes is not None and layer_num in deepstack_visual_indexes: - hidden_states_list.append(hidden_states) - - hidden_states = self.merger(hidden_states) - - # processing deepstack - if deepstack_visual_indexes is not None: - processed_hidden_states_list = [hidden_states] - for idx, x in enumerate(hidden_states_list): - x = self.merger_list[idx](x) - processed_hidden_states_list.append(x) - # we cat the original visual features and deepstack features - # along the feature dim - hidden_states = torch.cat( - processed_hidden_states_list, dim=1 - ) # [seq_len, hidden_size * (1 + depth_of_deepstack)] - - return hidden_states - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: - stacked_params_mapping = [ - # (param_name, shard_name, shard_id) - ("attn.qkv.", "attn.q.", "q"), - ("attn.qkv.", "attn.k.", "k"), - ("attn.qkv.", "attn.v.", "v"), - ] - params_dict = dict(self.named_parameters(remove_duplicate=False)) - loaded_params: set[str] = set() - - for name, loaded_weight in weights: - if name.endswith("patch_embed.proj.weight"): - loaded_weight = conv3d_to_linear_weight(loaded_weight) - - for param_name, weight_name, shard_id in stacked_params_mapping: - if weight_name not in name: - continue - name = name.replace(weight_name, param_name) - - param = params_dict[name] - weight_loader = param.weight_loader - weight_loader(param, loaded_weight, shard_id) - break - else: - param = params_dict[name] - weight_loader = getattr(param, "weight_loader", default_weight_loader) - weight_loader(param, loaded_weight) - loaded_params.add(name) - return loaded_params - - @support_torch_compile( dynamic_arg_dims={ "input_ids": 0, @@ -607,44 +113,7 @@ class Qwen3Omni_VisionTransformer(nn.Module): "deepstack_input_embeds": 0, } ) -class Qwen3MoeLLMModel(Qwen3MoeModel): - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__(vllm_config=vllm_config, prefix=prefix) - # from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import Qwen3OmniMoeTextRMSNorm - # config = vllm_config.model_config.hf_config - # self.norm = Qwen3OmniMoeTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) - self.deepstack_multiscale_layer_start = 1 - - def _deepstack_process(self, hidden_states: torch.Tensor, visual_embeds: torch.Tensor) -> torch.Tensor: - """ - Process deepstack visual embeddings by selectively adding them only to visual positions. - This matches the transformers implementation which uses visual_pos_masks. - - Args: - hidden_states: Hidden states tensor [num_tokens, hidden_size] - visual_embeds: Visual embeddings tensor [num_tokens, hidden_size] (zero-padded for non-visual positions) - - Returns: - Modified hidden_states with visual embeddings added only to visual positions - """ - # Derive visual position mask from visual_embeds (non-zero positions are visual) - # Use a small threshold to handle floating point precision issues - visual_pos_masks = torch.any(torch.abs(visual_embeds) > 1e-8, dim=-1) - - if visual_pos_masks.any(): - visual_pos_masks = visual_pos_masks.to(hidden_states.device) - visual_embeds = visual_embeds.to(hidden_states.device, hidden_states.dtype) - # Match transformers implementation EXACTLY (lines 2966-2969): - # In transformers: local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds - # Note: visual_embeds in transformers is NOT indexed, meaning it's already shaped - # [num_visual_tokens, hidden_size]. In vllm-omni, visual_embeds is [num_tokens, hidden_size] - # zero-padded, so we extract only visual positions to match transformers behavior. - visual_embeds_only = visual_embeds[visual_pos_masks, :] # Extract only visual positions - local_this = hidden_states[visual_pos_masks, :].clone() + visual_embeds_only - hidden_states[visual_pos_masks, :] = local_this - - return hidden_states - +class Qwen3MoeLLMModel(_Qwen3MoeLLMModel): def forward( self, input_ids: torch.Tensor, @@ -660,7 +129,7 @@ class Qwen3MoeLLMModel(Qwen3MoeModel): if inputs_embeds is not None: hidden_states = inputs_embeds else: - hidden_states = self.get_input_embeddings(input_ids) + hidden_states = self.embed_input_ids(input_ids) residual = None else: assert intermediate_tensors is not None @@ -683,11 +152,7 @@ class Qwen3MoeLLMModel(Qwen3MoeModel): ) if deepstack_input_embeds is not None and layer_idx in range(0, len(deepstack_input_embeds)): - # Use _deepstack_process to match transformers implementation - # This selectively adds visual embeddings only to visual positions - hidden_states = self._deepstack_process( - hidden_states, deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"] - ) + hidden_states = hidden_states + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"] if not get_pp_group().is_last_rank: return IntermediateTensors({"hidden_states": hidden_states, "residual": residual}) @@ -856,7 +321,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( # Derive audio placeholders from video placeholders mm_placeholders = self._derive_audio_from_video_placeholders(mm_placeholders, mm_prompt_updates) else: - prompt_ids, prompt, mm_placeholders = self._apply_prompt_updates( + prompt_ids, mm_placeholders = self._apply_prompt_updates( prompt_ids, mm_prompt_updates, ) @@ -866,7 +331,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( mm_item_counts, ) - return prompt_ids, prompt, mm_placeholders + return prompt_ids, mm_placeholders def get_updates_use_audio_in_video( self, @@ -962,7 +427,6 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( def get_replacement_qwen2_vision(item_idx: int, modality: str): grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx] - # grid_thw = out_mm_kwargs[f"{modality}_grid_thw"][item_idx] assert isinstance(grid_thw, torch.Tensor) merge_length = image_processor.merge_size**2 @@ -974,8 +438,6 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( def get_replacement_qwen2_use_audio_in_video(item_idx: int): nonlocal audio_in_video_item_idx - # audio_num_features = audio_output_lengths[audio_item_idx + item_idx] - # video_grid_thw = out_mm_kwargs["video_grid_thw"][item_idx] audio_num_features = audio_output_lengths[audio_in_video_item_idx + item_idx] video_grid_thw = out_mm_data["video_grid_thw"][item_idx] @@ -1022,7 +484,7 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( def _derive_audio_from_video_placeholders( self, placeholders: Mapping[str, list[PlaceholderFeaturesInfo]], - mm_prompt_updates: Mapping[str, list[PromptUpdateDetails]], + mm_prompt_updates: MultiModalPromptUpdates, ) -> Mapping[str, list[PlaceholderFeaturesInfo]]: """ Helper to derive audio placeholders from video placeholders when @@ -1116,8 +578,6 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix input_features = audio_input["input_features"] audio_feature_lengths = audio_input["audio_feature_lengths"] - audio_feature_lengths = flatten_bn(audio_feature_lengths, concat=True) - audio_feat_lengths, audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths) audio_outputs = self.audio_tower( @@ -1138,6 +598,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( nn.Module, SupportsMultiModal, SupportsPP, + SupportsMRoPE, Qwen3OmniMoeConditionalGenerationMixin, ): merge_by_field_config = True @@ -1163,6 +624,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): super().__init__() + self.vllm_config = vllm_config # needed for torch compile forward context thinker_config: Qwen3OmniMoeThinkerConfig = vllm_config.model_config.hf_config quant_config = vllm_config.quant_config multimodal_config = vllm_config.model_config.multimodal_config @@ -1184,24 +646,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config) - # attn_backend_override = ( - # multimodal_config.mm_encoder_attn_backend - # if multimodal_config is not None - # else None - # ) + attn_backend_override = multimodal_config.mm_encoder_attn_backend if multimodal_config is not None else None self.visual = Qwen3Omni_VisionTransformer( vision_config=thinker_config.vision_config, norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), quant_config=quant_config, prefix=maybe_prefix(prefix, "visual"), - # attn_backend_override=attn_backend_override, + attn_backend_override=attn_backend_override, ) self.quant_config = quant_config - lm_vllm_config = vllm_config.with_hf_config(thinker_config.text_config, architectures=["Qwen3MoeForCausalLM"]) - self.language_model = Qwen3MoeLLMForCausalLM( - vllm_config=lm_vllm_config, + vllm_config=vllm_config.with_hf_config(thinker_config.text_config, architectures=["Qwen3MoeForCausalLM"]), prefix=maybe_prefix(prefix, "language_model"), ) @@ -1274,7 +730,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( def get_language_model(self) -> torch.nn.Module: return self.language_model - def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings | None: + def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings | None: mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs) if not mm_input_by_modality: return [] @@ -1298,27 +754,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( multimodal_embeddings += tuple(audio_embeddings) return multimodal_embeddings - def _get_text_embeddings( - self, - input_ids: Tensor, - embed_input_ids: Callable[[Tensor], Tensor], - *, - is_multimodal: Tensor | None, - handle_oov_mm_token: bool, - ) -> Tensor: - if handle_oov_mm_token and is_multimodal is not None: - is_text = ~is_multimodal - text_embeds = embed_input_ids(input_ids[is_text]) - - return torch.empty( - (input_ids.shape[0], text_embeds.shape[1]), - dtype=text_embeds.dtype, - device=text_embeds.device, - ).masked_scatter_(is_text.unsqueeze_(-1), text_embeds) - - return embed_input_ids(input_ids) - - def get_input_embeddings( + def embed_input_ids( self, input_ids: torch.Tensor, multimodal_embeddings: MultiModalEmbeddings | None = None, @@ -1326,18 +762,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( is_multimodal: torch.Tensor | None = None, handle_oov_mm_token: bool = False, ) -> torch.Tensor: - # Compute is_multimodal if not provided and we have multimodal embeddings - if is_multimodal is None and multimodal_embeddings is not None and len(multimodal_embeddings) > 0: - # Create a boolean tensor marking which positions are multimodal tokens - is_multimodal = ( - (input_ids == self.config.image_token_id) - | (input_ids == self.config.video_token_id) - | (input_ids == self.config.audio_token_id) - ) - - inputs_embeds = self._get_text_embeddings( + inputs_embeds = self._embed_text_input_ids( input_ids, - self.language_model.get_input_embeddings, + self.language_model.embed_input_ids, is_multimodal=is_multimodal, handle_oov_mm_token=handle_oov_mm_token, ) @@ -1406,7 +833,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( positions: torch.Tensor, intermediate_tensors: IntermediateTensors | None = None, inputs_embeds: torch.Tensor | None = None, - *, capture_layer_indices: Sequence[int] | None = None, return_hidden_states: bool = False, **kwargs: object, @@ -1415,7 +841,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( inputs_embeds = None if self.use_deepstack and inputs_embeds is not None and get_pp_group().is_first_rank: - deepstack_input_embeds = self._get_deepstack_input_embeds(input_ids.size(1)) + deepstack_input_embeds = self._get_deepstack_input_embeds(inputs_embeds.size(0)) else: deepstack_input_embeds = None @@ -1431,14 +857,13 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ) if inputs_embeds is not None and get_pp_group().is_first_rank: - self._clear_deepstack_input_embeds(input_ids.size(1)) + self._clear_deepstack_input_embeds(inputs_embeds.size(0)) - return hidden_states.unsqueeze(0), captured_hidden_states + return hidden_states, captured_hidden_states def compute_logits( self, hidden_states: torch.Tensor, - sampling_metadata: SamplingMetadata = None, ) -> torch.Tensor | None: return self.language_model.compute_logits(hidden_states) @@ -1449,55 +874,48 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( ) loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper) - # Log load summary - try: - total_bytes = 0 - for name, param in self.named_parameters(): - if param is not None and param.data is not None: - total_bytes += param.data.numel() * param.data.element_size() - device = next(self.parameters()).device - logger.info( - "[Model Loaded] name=%s, success=%s, size=%.2f MB, device=%s", - self.__class__.__name__, - True, - total_bytes / (1024**2), - str(device), - ) - except Exception: - pass - return loaded_weights def get_mrope_input_positions( self, input_tokens: list[int], - hf_config: PretrainedConfig, - image_grid_thw: list[list[int]] | torch.Tensor | None, - video_grid_thw: list[list[int]] | torch.Tensor | None, - second_per_grid_ts: list[float] | None = None, - context_len: int = 0, - seq_len: int | None = None, - audio_feature_lengths: torch.Tensor | None = None, - use_audio_in_video: bool = False, + mm_features: list[MultiModalFeatureSpec], ) -> tuple[torch.Tensor, int]: - config = hf_config.thinker_config - if isinstance(image_grid_thw, list): - image_grid_thw = torch.tensor(image_grid_thw) - if isinstance(video_grid_thw, list): - video_grid_thw = torch.tensor(video_grid_thw) + kwargs = MultiModalFeatureSpec.gather_kwargs( + mm_features, + { + "image_grid_thw", + "video_grid_thw", + "second_per_grid_ts", + "audio_feature_lengths", + "use_audio_in_video", + }, + ) + image_grid_thw = kwargs.get("image_grid_thw", []) + video_grid_thw = kwargs.get("video_grid_thw", []) + second_per_grid_ts = kwargs.get("second_per_grid_ts", []) + audio_feature_lengths = kwargs.get("audio_feature_lengths", []) + use_audio_in_video = any(kwargs.get("use_audio_in_video", [])) + + image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(image_grid_thw) + video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)(video_grid_thw) + input_ids = torch.tensor(input_tokens) if input_ids is None or input_ids.ndim != 1: raise ValueError("_omni3_get_input_positions_tensor expects 1D input_ids") seq_len = input_ids.shape[0] - if audio_feature_lengths is not None and not isinstance(audio_feature_lengths, torch.Tensor): - audio_feature_lengths = torch.as_tensor(audio_feature_lengths, dtype=torch.long) + + if isinstance(audio_feature_lengths, list): + audio_feature_lengths = torch.tensor(audio_feature_lengths, dtype=torch.long) + if not len(second_per_grid_ts) and len(video_grid_thw): second_per_grid_ts = 2.0 second_per_grids = torch.ones(len(video_grid_thw), dtype=torch.float32) * second_per_grid_ts else: second_per_grids = torch.tensor(second_per_grid_ts, dtype=torch.float32) + config = self.config spatial_merge_size = config.vision_config.spatial_merge_size image_token_id = config.image_token_id video_token_id = config.video_token_id diff --git a/vllm_omni/model_executor/models/registry.py b/vllm_omni/model_executor/models/registry.py index fd64b8b2e..56bceae41 100644 --- a/vllm_omni/model_executor/models/registry.py +++ b/vllm_omni/model_executor/models/registry.py @@ -10,12 +10,12 @@ _OMNI_MODELS = { "qwen2_5_omni", "qwen2_5_omni_thinker", "Qwen2_5OmniThinkerForConditionalGeneration", - ), # noqa: E501 + ), "Qwen2_5OmniTalkerModel": ( "qwen2_5_omni", "qwen2_5_omni_talker", "Qwen2_5OmniTalkerForConditionalGeneration", - ), # noqa: E501 + ), "Qwen2_5OmniToken2WavModel": ( "qwen2_5_omni", "qwen2_5_omni_token2wav", diff --git a/vllm_omni/request.py b/vllm_omni/request.py index 6190fad00..ca5bd437d 100644 --- a/vllm_omni/request.py +++ b/vllm_omni/request.py @@ -2,7 +2,6 @@ from collections.abc import Callable from typing import TYPE_CHECKING from vllm.v1.request import Request -from vllm.v1.structured_output.request import StructuredOutputRequest if TYPE_CHECKING: from vllm.v1.core.kv_cache_utils import BlockHash @@ -57,19 +56,16 @@ class OmniRequest(Request): request_id=request.request_id, client_index=request.client_index, prompt_token_ids=request.prompt_token_ids, + prompt_embeds=request.prompt_embeds, mm_features=request.mm_features, sampling_params=request.sampling_params, pooling_params=request.pooling_params, eos_token_id=request.eos_token_id, 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 - ), cache_salt=request.cache_salt, priority=request.priority, trace_headers=request.trace_headers, block_hasher=block_hasher, - prompt_embeds=request.prompt_embeds, additional_information=request.additional_information, ) diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 48cdd29f9..162c66dd4 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -6,11 +6,17 @@ and also outputs sampled tokens. from __future__ import annotations +from copy import copy +from typing import Any, NamedTuple + +import numpy as np import torch -from vllm.forward_context import BatchDescriptor +from vllm.config import CUDAGraphMode +from vllm.forward_context import set_forward_context from vllm.logger import init_logger -from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput from vllm.v1.outputs import AsyncModelRunnerOutput +from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( @@ -20,7 +26,6 @@ from vllm.v1.worker.gpu_model_runner import ( get_pp_group, get_tp_group, has_kv_transfer_group, - set_forward_context, ) from vllm.v1.worker.utils import is_residual_scattered_for_sp @@ -30,14 +35,25 @@ from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner logger = init_logger(__name__) +class ExecuteModelState(NamedTuple): + scheduler_output: SchedulerOutput + logits: torch.Tensor | None + spec_decode_metadata: Any + spec_decode_common_attn_metadata: Any + hidden_states: torch.Tensor + sample_hidden_states: torch.Tensor + aux_hidden_states: list[torch.Tensor] | None + ec_connector_output: Any + multimodal_outputs: Any + + class GPUARModelRunner(OmniGPUModelRunner): """Autoregressive GPU model runner that returns hidden states per request. - This runner follows the same preparation and forward path as GPUModelRunner - (inputs assembly, multi-modal handling, TP/PP/DP integration, CUDA graphs), - and additionally performs lightweight sampling so that sampled tokens are - available in outputs. Hidden representations are taken at the same indices - that GPUModelRunner would use for sampling/logits (i.e. `logits_indices`). + Follows the v0.12 two-phase execute/sample flow from GPUModelRunner, and + reuses Omni hooks for additional_information / multimodal outputs. This + class only overrides sample_tokens to expose hidden states + multimodal + outputs per request while keeping Async output semantics. """ def __init__(self, *args, **kwargs): @@ -65,19 +81,14 @@ class GPUARModelRunner(OmniGPUModelRunner): self, scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, - ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None: with record_function_or_nullcontext("Preprocess"): with self.synchronize_input_prep(): - # Update persistent batch states. self._update_states(scheduler_output) - - # Decode per-request prompt_embeds / additional_information payloads - # (if present) into CPU tensors self._decode_and_store_request_payloads(scheduler_output) 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) if self.cache_config.kv_sharing_fast_prefill: @@ -87,64 +98,83 @@ class GPUARModelRunner(OmniGPUModelRunner): "it when the requests need prompt logprobs" ) - # Prepare the decoder inputs. + num_reqs = self.input_batch.num_reqs + req_ids = self.input_batch.req_ids + tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] + num_scheduled_tokens_np = np.array(tokens, dtype=np.int32) + max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) + num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens + + logits_indices, spec_decode_metadata = self._prepare_inputs( + scheduler_output, + num_scheduled_tokens_np, + ) + + ( + cudagraph_mode, + batch_desc, + ubatch_slices, + num_tokens_across_dp, + ) = self._determine_batch_execution_and_padding( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs, + num_scheduled_tokens_np=num_scheduled_tokens_np, + max_num_scheduled_tokens=max_num_scheduled_tokens, + use_cascade_attn=False, + ) + + num_tokens_padded = batch_desc.num_tokens + num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + pad_attn = cudagraph_mode == CUDAGraphMode.FULL + ( attn_metadata, - logits_indices, - spec_decode_metadata, - num_scheduled_tokens_np, spec_decode_common_attn_metadata, - max_query_len, - ubatch_slices, - num_tokens_after_padding, - ) = self._prepare_inputs(scheduler_output) + ) = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded if pad_attn else None, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded if pad_attn else None, + max_query_len=max_num_scheduled_tokens, + ubatch_slices=ubatch_slices, + logits_indices=logits_indices, + use_spec_decode=use_spec_decode, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + cascade_attn_prefix_lens=None, + ) ( - num_scheduled_tokens, - num_input_tokens, - num_tokens_across_dp, input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, - per_req_additional_information, + ec_connector_output, ) = self._preprocess( scheduler_output, - num_scheduled_tokens_np, + num_tokens_padded, intermediate_tensors, - ubatch_slices, - num_tokens_after_padding, ) - 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) + if self.calculate_kv_scales: + cudagraph_mode = CUDAGraphMode.NONE + self.calculate_kv_scales = False - # This is currently to get around the assert in the DPMetadata - # where it wants `num_tokens_across_dp` to align with `num_tokens` - if ubatch_slices is not None: - num_input_tokens = ubatch_slices[0].num_tokens - - # Run the model. - # Use persistent buffers for CUDA graphs. with ( set_forward_context( attn_metadata, self.vllm_config, - num_tokens=num_input_tokens, + num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, + cudagraph_runtime_mode=cudagraph_mode, + batch_descriptor=batch_desc, ubatch_slices=ubatch_slices, ), record_function_or_nullcontext("Forward"), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, ): - model_kwargs_extra = self._build_model_kwargs_extra(per_req_additional_information, num_scheduled_tokens_np) - model_output = self.model( + model_output = self._model_forward( input_ids=input_ids, positions=positions, intermediate_tensors=intermediate_tensors, @@ -153,10 +183,9 @@ class GPUARModelRunner(OmniGPUModelRunner): sampling_metadata=self.input_batch.sampling_metadata, logits_index=logits_indices, sampler=self.sampler, - **model_kwargs_extra, ) - with record_function_or_nullcontext("Postprocess"): + with record_function_or_nullcontext("gpu_model_runner: postprocess"): if self.use_aux_hidden_state_outputs: # True when EAGLE 3 is used. hidden_states, aux_hidden_states = model_output @@ -165,38 +194,42 @@ class GPUARModelRunner(OmniGPUModelRunner): hidden_states = model_output aux_hidden_states = None - # hidden_states, multimodal_outputs = self.extract_multimodal_outputs( - # hidden_states - # ) multimodal_outputs = model_output.multimodal_outputs hidden_states = model_output.text_hidden_states - # The model side may return per-request additional_information updates (model-agnostic channel). - # Convention: multimodal_outputs["additional_information_update"] is a list[dict] in batch order; - # the runner merges it into the corresponding request's additional_information_cpu for subsequent decode. - self._process_additional_information_updates(multimodal_outputs) + + if multimodal_outputs is not None: + keys_or_type = ( + list(multimodal_outputs.keys()) + if isinstance(multimodal_outputs, dict) + else type(multimodal_outputs) + ) + logger.debug(f"[AR] execute_model: multimodal_outputs keys = {keys_or_type}") + else: + logger.debug("[AR] execute_model: multimodal_outputs is None") + if not self.broadcast_pp_output: - # Common case. if not get_pp_group().is_last_rank: - # Return the intermediate tensors. assert isinstance(hidden_states, IntermediateTensors) hidden_states.kv_connector_output = kv_connector_output return hidden_states 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_tokens_padded, + num_scheduled_tokens_np, + ) output.kv_connector_output = kv_connector_output return output sample_hidden_states = hidden_states[logits_indices] logits = self.model.compute_logits(sample_hidden_states) else: - # Rare case. assert not self.is_pooling_model 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_tokens_padded) } get_pp_group().send_tensor_dict( hidden_states.tensors, @@ -208,26 +241,71 @@ class GPUARModelRunner(OmniGPUModelRunner): sample_hidden_states = hidden_states[logits_indices] logits = self.model.compute_logits(sample_hidden_states) - model_output_broadcast_data = {} + model_output_broadcast_data: dict[str, Any] = {} if logits is not None: model_output_broadcast_data["logits"] = logits.contiguous() - model_output_broadcast_data = get_pp_group().broadcast_tensor_dict( + broadcasted = get_pp_group().broadcast_tensor_dict( model_output_broadcast_data, src=len(get_pp_group().ranks) - 1 ) - assert model_output_broadcast_data is not None - logits = model_output_broadcast_data["logits"] + assert broadcasted is not None + logits = broadcasted["logits"] - # 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) + self.execute_model_state = ExecuteModelState( + scheduler_output, + logits, + spec_decode_metadata, + spec_decode_common_attn_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + ec_connector_output, + multimodal_outputs, + ) + self.kv_connector_output = kv_connector_output + return None - with record_function_or_nullcontext("Sample"): + @torch.inference_mode() + def sample_tokens( + self, + grammar_output: GrammarOutput | None, + ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None + + if self.execute_model_state is None: + if not kv_connector_output: + return None # type: ignore[return-value] + if kv_connector_output.is_empty(): + return EMPTY_MODEL_RUNNER_OUTPUT + output = copy(EMPTY_MODEL_RUNNER_OUTPUT) + output.kv_connector_output = kv_connector_output + return output + + ( + scheduler_output, + logits, + spec_decode_metadata, + spec_decode_common_attn_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + ec_connector_output, + multimodal_outputs, + ) = self.execute_model_state + self.execute_model_state = None + + if grammar_output is not None: + apply_grammar_bitmask(scheduler_output, grammar_output, self.input_batch, logits) + + with record_function_or_nullcontext("gpu_model_runner: sample"): sampler_output = self._sample(logits, spec_decode_metadata) + self.input_batch.prev_sampled_token_ids = None + def propose_draft_token_ids(sampled_token_ids): assert spec_decode_common_attn_metadata is not None - with record_function_or_nullcontext("Draft"): + with record_function_or_nullcontext("gpu_model_runner: draft"): self._draft_token_ids = self.propose_draft_token_ids( scheduler_output, sampled_token_ids, @@ -239,30 +317,40 @@ class GPUARModelRunner(OmniGPUModelRunner): spec_decode_common_attn_metadata, ) + spec_config = self.speculative_config use_padded_batch_for_eagle = ( - self.speculative_config - and self.speculative_config.use_eagle() - and not self.speculative_config.disable_padded_drafter_batch + spec_config is not None and spec_config.use_eagle() and not spec_config.disable_padded_drafter_batch ) effective_drafter_max_model_len = self.max_model_len if effective_drafter_max_model_len is None: effective_drafter_max_model_len = self.model_config.max_model_len if ( - self.speculative_config - and self.speculative_config.draft_model_config is not None - and self.speculative_config.draft_model_config.max_model_len is not None + spec_config is not None + and spec_config.draft_model_config is not None + and spec_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 = spec_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 - <= effective_drafter_max_model_len + spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens <= effective_drafter_max_model_len ) - if use_padded_batch_for_eagle and input_fits_in_drafter: - # EAGLE speculative decoding can use the GPU sampled tokens - # as inputs, and does not need to wait for bookkeeping to finish. - propose_draft_token_ids(sampler_output.sampled_token_ids) + if use_padded_batch_for_eagle: + assert self.speculative_config is not None + assert isinstance(self.drafter, EagleProposer) + sampled_token_ids = sampler_output.sampled_token_ids + if input_fits_in_drafter: + propose_draft_token_ids(sampled_token_ids) + elif self.valid_sampled_token_count_event is not None: + assert spec_decode_common_attn_metadata is not None + next_token_ids, valid_sampled_tokens_count = self.drafter.prepare_next_token_ids_padded( + spec_decode_common_attn_metadata, + sampled_token_ids, + self.requests, + self.input_batch, + self.discard_request_mask.gpu, + ) + self._copy_valid_sampled_token_count(next_token_ids, valid_sampled_tokens_count) - with record_function_or_nullcontext("Bookkeep"): + with record_function_or_nullcontext("gpu_model_runner: bookkeep"): ( num_nans_in_logits, logprobs_lists, @@ -276,82 +364,91 @@ class GPUARModelRunner(OmniGPUModelRunner): sampler_output, logits, hidden_states, - num_scheduled_tokens, + scheduler_output.total_num_scheduled_tokens, + spec_decode_metadata, ) 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) - with record_function_or_nullcontext("EPLB"): + with record_function_or_nullcontext("gpu_model_runner: eplb"): self.eplb_step() - # Convert to per-request tensors on CPU + self._process_additional_information_updates(multimodal_outputs) + hidden_states_cpu = hidden_states.detach().to("cpu").contiguous() - # pooler_output: list[torch.Tensor | None] = [] - pooler_output: list[torch.Tensor | None] = [] - prev_logits_index = 0 - for rid, logits_index in zip(req_ids_output_copy, logits_indices): - # Base payload: hidden slice for this request in this iteration - hidden_slice = hidden_states_cpu[prev_logits_index : logits_index + 1] + num_scheduled_tokens_np = getattr(self, "_omni_num_scheduled_tokens_np", None) + if num_scheduled_tokens_np is None: + req_ids = self.input_batch.req_ids + num_scheduled_tokens_np = np.array( + [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids], + dtype=np.int32, + ) + + pooler_output: list[dict[str, object]] = [] + for rid in req_ids_output_copy: + idx = req_id_to_index_output_copy[rid] + start = int(self.query_start_loc.cpu[idx]) + sched = int(num_scheduled_tokens_np[idx]) + end = start + sched + hidden_slice = hidden_states_cpu[start:end] payload: dict[str, object] = {"hidden": hidden_slice} - # Merge multimodal_outputs if present if isinstance(multimodal_outputs, dict) and multimodal_outputs: mm_payload: dict[str, object] = {} for k, v in multimodal_outputs.items(): try: - # Case 1: tensor aligned on token dimension if isinstance(v, torch.Tensor) and v.shape[0] == hidden_states_cpu.shape[0]: - mm_payload[k] = v.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous() - elif isinstance(v, torch.Tensor) and v.shape[0] != hidden_states_cpu.shape[0]: - logger.error( - f"Error in merge multimodal outputs: Tensor dimension mismatch, \ - {v.shape} != {hidden_states_cpu.shape} for {k}" - ) - # Case 2: nested dict of tensors aligned on token dimension (e.g., selected_hidden_layers) + mm_payload[k] = v.detach().to("cpu")[start:end].contiguous() elif isinstance(v, dict): sub_dict: dict[str, torch.Tensor] = {} for sk, sv in v.items(): if isinstance(sv, torch.Tensor) and sv.shape[0] == hidden_states_cpu.shape[0]: - sub_dict[str(sk)] = ( - sv.detach().to("cpu")[prev_logits_index : logits_index + 1].contiguous() - ) + sub_dict[str(sk)] = sv.detach().to("cpu")[start:end].contiguous() if sub_dict: mm_payload[k] = sub_dict elif isinstance(v, list): element = v[0] if isinstance(element, torch.Tensor): element = element.detach().to("cpu").contiguous() - multimodal_outputs[k] = v[1:] if len(v) > 1 else v mm_payload[k] = element except Exception as e: - # Best-effort; skip malformed entries logger.error(f"Error in merge multimodal outputs: {e}") if mm_payload: payload.update(mm_payload) - pooler_output.append(payload) # type: ignore[arg-type] - prev_logits_index = logits_index + 1 - output = OmniModelRunnerOutput( - req_ids=req_ids_output_copy, - req_id_to_index=req_id_to_index_output_copy, - 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), - kv_connector_output=kv_connector_output, - num_nans_in_logits=num_nans_in_logits, - ) + pooler_output.append(payload) + with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): + output = OmniModelRunnerOutput( + req_ids=req_ids_output_copy, + req_id_to_index=req_id_to_index_output_copy, + 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), + kv_connector_output=kv_connector_output, + ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, + num_nans_in_logits=num_nans_in_logits, + ) if not self.use_async_scheduling: return output + with record_function_or_nullcontext("gpu_model_runner: AsyncGPUModelRunnerOutput"): + async_output = AsyncGPUModelRunnerOutput( + model_runner_output=output, + sampled_token_ids=sampler_output.sampled_token_ids, + logprobs_tensors=sampler_output.logprobs_tensors, + invalid_req_indices=invalid_req_indices, + async_output_copy_stream=self.async_output_copy_stream, + vocab_size=self.input_batch.vocab_size, + ) + with record_function_or_nullcontext("gpu_model_runner: set_async_sampled_token_ids"): + # Save ref of sampled_token_ids CPU tensor if the batch contains + # any requests with sampling params that require output ids. + self.input_batch.set_async_sampled_token_ids( + async_output.sampled_token_ids_cpu, + async_output.async_copy_ready_event, + ) - return AsyncGPUModelRunnerOutput( - model_runner_output=output, - sampled_token_ids=sampler_output.sampled_token_ids, - invalid_req_indices=invalid_req_indices, - async_output_copy_stream=self.async_output_copy_stream, - ) + return async_output def _merge_additional_information_update(self, req_id: str, upd: dict) -> None: req_state = self.requests.get(req_id) @@ -365,13 +462,9 @@ class GPUARModelRunner(OmniGPUModelRunner): if isinstance(v, torch.Tensor): merged[k] = v.detach().to("cpu").contiguous() elif isinstance(v, list): - new_list = [] - for item in v: - if isinstance(item, torch.Tensor): - new_list.append(item.detach().to("cpu").contiguous()) - else: - new_list.append(item) - merged[k] = new_list + merged[k] = [ + (item.detach().to("cpu").contiguous() if isinstance(item, torch.Tensor) else item) for item in v + ] else: merged[k] = v setattr(req_state, "additional_information_cpu", merged) diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 0c34e1ad6..4748d1fd3 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -4,7 +4,8 @@ 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.utils.mem_constants import GiB_bytes +from vllm.utils.mem_utils import 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 @@ -20,19 +21,37 @@ class GPUARWorker(GPUWorker): """ def init_device(self): - """Initialize the GPU device and autoregressive model runner. - - Sets up CUDA device, initializes distributed environment, - sets random seed, and creates a GPUARModelRunner instance. - - Raises: - ValueError: If free GPU memory is insufficient for the - requested memory utilization - RuntimeError: If device type is not supported - """ - if self.device_config.device.type == "cuda": + device = self.device_config.device + if isinstance(device, torch.device) and device.type == "cuda": # This env var set by Ray causes exceptions with graph building. os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + if ( + self.parallel_config.data_parallel_size > 1 + and self.parallel_config.data_parallel_size_local > 0 + and self.parallel_config.distributed_executor_backend not in ["ray", "external_launcher"] + and self.vllm_config.parallel_config.data_parallel_backend != "ray" + and self.vllm_config.parallel_config.nnodes_within_dp == 1 + ): + # Use local DP rank if available, otherwise use global DP rank. + dp_local_rank = self.parallel_config.data_parallel_rank_local + if dp_local_rank is None: + dp_local_rank = self.parallel_config.data_parallel_rank + + tp_pp_world_size = ( + self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size + ) + + # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK + self.local_rank += dp_local_rank * tp_pp_world_size + assert self.local_rank < torch.cuda.device_count(), ( + f"DP adjusted local rank {self.local_rank} is out of bounds. " + ) + visible_device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0 + assert self.parallel_config.local_world_size <= visible_device_count, ( + f"local_world_size ({self.parallel_config.local_world_size}) must " + f"be less than or equal to the number of visible devices " + f"({visible_device_count})." + ) self.device = torch.device(f"cuda:{self.local_rank}") current_platform.set_device(self.device) @@ -62,23 +81,23 @@ class GPUARWorker(GPUWorker): self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization if self.init_snapshot.free_memory < self.requested_memory: - def GiB(b): - return round(b / GiB_bytes, 2) + def gib(bytes_val: float) -> float: + return round(bytes_val / GiB_bytes, 2) raise ValueError( f"Free memory on device " - f"({GiB(self.init_snapshot.free_memory)}/" - f"{GiB(self.init_snapshot.total_memory)} GiB) on startup " + f"({gib(self.init_snapshot.free_memory)}/" + f"{gib(self.init_snapshot.total_memory)} GiB) on startup " f"is less than desired GPU memory utilization " f"({self.cache_config.gpu_memory_utilization}, " - f"{GiB(self.requested_memory)} GiB). Decrease GPU memory " + f"{gib(self.requested_memory)} GiB). Decrease GPU memory " f"utilization or reduce GPU memory used by other processes." ) else: raise RuntimeError(f"Not support device type: {self.device_config.device}") # Construct the model runner - self.model_runner: GPUARModelRunner = GPUARModelRunner(self.vllm_config, self.device) + self.model_runner = 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_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index fa5cab9e6..c5bffa56b 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -8,14 +8,12 @@ from __future__ import annotations import gc import logging -from typing import Any import numpy as np import torch from vllm.config import CUDAGraphMode -from vllm.forward_context import BatchDescriptor from vllm.multimodal.inputs import MultiModalKwargs -from vllm.v1.attention.backends.utils import CommonAttentionMetadata +from vllm.utils.math_utils import cdiv from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.utils import record_function_or_nullcontext @@ -27,7 +25,6 @@ from vllm.v1.worker.gpu_model_runner import ( get_pp_group, set_forward_context, ) -from vllm.v1.worker.ubatch_utils import UBatchSlices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs from vllm_omni.outputs import OmniModelRunnerOutput @@ -44,110 +41,6 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): - Executes generation process and returns tensors via `pooler_output`. """ - def _preprocess( - self, - scheduler_output: SchedulerOutput, - intermediate_tensors: IntermediateTensors | None = None, - ubatch_slices: UBatchSlices | None = None, - num_tokens_after_padding: torch.Tensor | None = None, - ) -> tuple[ - int, - int, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor, - IntermediateTensors | None, - dict[str, Any], - ]: - # Input token count for this iteration (not used by diffusion, but - # retained to keep DP padding/ordering consistent) - num_input_tokens = scheduler_output.total_num_scheduled_tokens - num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens) - num_input_tokens += num_pad - - # _prepare_inputs may reorder the batch, so we must gather multi - # modal outputs after that to ensure the correct order - if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder: - # Run the multimodal encoder if any. - self._execute_mm_encoder(scheduler_output) - mm_embeds = self._gather_mm_embeddings(scheduler_output) - - # NOTE(woosuk): To unify token ids and soft tokens (vision - # embeddings), we always use embeddings (rather than token ids) - # as input to the multimodal model, even when the input is text. - inputs_embeds_scheduled = self.model.get_input_embeddings( - input_ids=self.input_ids.gpu[:num_input_tokens], - multimodal_embeddings=mm_embeds or None, - ) - - # TODO(woosuk): Avoid the copy. Optimize. - self.inputs_embeds.gpu[:num_input_tokens].copy_(inputs_embeds_scheduled) - - input_ids = self.input_ids.gpu[:num_input_tokens] - inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] - model_kwargs = { - **self._init_model_kwargs(num_input_tokens), - **self._extract_mm_kwargs(scheduler_output), - } - elif self.enable_prompt_embeds and get_pp_group().is_first_rank: - # Get the input embeddings for the tokens that are not input embeds, - # then put them into the appropriate positions. - # TODO(qthequartermasterman): Since even when prompt embeds are - # enabled, (a) not all requests will use prompt embeds, and (b) - # after the initial prompt is processed, the rest of the generated - # tokens will be token ids, it is not desirable to have the - # embedding layer outside of the CUDA graph all the time. The v0 - # engine avoids this by "double compiling" the CUDA graph, once - # with input_ids and again with inputs_embeds, for all num_tokens. - # If a batch only has token ids, then including the embedding layer - # in the CUDA graph will be more performant (like in the else case - # below). - token_ids_idx = self.is_token_ids.gpu[:num_input_tokens].nonzero(as_tuple=False).squeeze(1) - # Some tokens ids may need to become embeds - if token_ids_idx.numel() > 0: - token_ids = self.input_ids.gpu[token_ids_idx] - tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids) - self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds - - inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] - model_kwargs = self._init_model_kwargs(num_input_tokens) - input_ids = self.input_ids.gpu[:num_input_tokens] - else: - # For text-only models, we use token ids as input. - # While it is possible to use embeddings as input just like the - # multimodal models, it is not desirable for performance since - # then the embedding layer is not included in the CUDA graph. - input_ids = self.input_ids.gpu[:num_input_tokens] - inputs_embeds = None - model_kwargs = self._init_model_kwargs(num_input_tokens) - if self.uses_mrope: - positions = self.mrope_positions.gpu[:, :num_input_tokens] - else: - positions = self.positions.gpu[:num_input_tokens] - - if get_pp_group().is_first_rank: - intermediate_tensors = None - else: - intermediate_tensors = self.sync_and_slice_intermediate_tensors( - num_input_tokens, intermediate_tensors, True - ) - - if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: - encoder_inputs = self._extract_encoder_inputs(scheduler_output) - model_kwargs.update(encoder_inputs) - - return ( - num_input_tokens, - num_input_tokens, - num_tokens_after_padding, - input_ids, - inputs_embeds, - positions, - intermediate_tensors, - model_kwargs, - ) - @torch.inference_mode() def execute_model( self, @@ -156,52 +49,81 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): ) -> OmniModelRunnerOutput | IntermediateTensors: with record_function_or_nullcontext("Preprocess"): with self.synchronize_input_prep(): - # Update persistent batch states. self._update_states(scheduler_output) if not scheduler_output.total_num_scheduled_tokens: return EMPTY_MODEL_RUNNER_OUTPUT - # Prepare the decoder inputs. + num_reqs = self.input_batch.num_reqs + req_ids = self.input_batch.req_ids + tokens = [scheduler_output.num_scheduled_tokens[i] for i in req_ids] + num_scheduled_tokens_np = np.array(tokens, dtype=np.int32) + max_num_scheduled_tokens = int(num_scheduled_tokens_np.max()) + num_tokens_unpadded = scheduler_output.total_num_scheduled_tokens + + logits_indices, spec_decode_metadata = self._prepare_inputs( + scheduler_output, + num_scheduled_tokens_np, + ) + + ( + cudagraph_mode, + batch_desc, + ubatch_slices, + num_tokens_across_dp, + ) = self._determine_batch_execution_and_padding( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs, + num_scheduled_tokens_np=num_scheduled_tokens_np, + max_num_scheduled_tokens=max_num_scheduled_tokens, + use_cascade_attn=False, + ) + + num_tokens_padded = batch_desc.num_tokens + num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + pad_attn = cudagraph_mode == CUDAGraphMode.FULL + ( attn_metadata, - logits_indices, - spec_decode_metadata, - num_scheduled_tokens_np, spec_decode_common_attn_metadata, - max_query_len, - ubatch_slices, - num_tokens_after_padding, - ) = self._prepare_inputs(scheduler_output) + ) = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_tokens_padded=num_tokens_padded if pad_attn else None, + num_reqs=num_reqs, + num_reqs_padded=num_reqs_padded if pad_attn else None, + max_query_len=max_num_scheduled_tokens, + ubatch_slices=ubatch_slices, + logits_indices=logits_indices, + use_spec_decode=use_spec_decode, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + cascade_attn_prefix_lens=None, + ) - # Input token count for this iteration (not used by diffusion, but - # retained to keep DP padding/ordering consistent) ( - num_scheduled_tokens, - num_input_tokens, - num_tokens_across_dp, input_ids, inputs_embeds, positions, intermediate_tensors, model_kwargs, + ec_connector_output, ) = self._preprocess( scheduler_output, + num_tokens_padded, intermediate_tensors, - ubatch_slices, - num_tokens_after_padding, ) - cudagraph_runtime_mode = CUDAGraphMode.NONE - # Run the model. - # Use persistent buffers for CUDA graphs. + if self.calculate_kv_scales: + cudagraph_mode = CUDAGraphMode.NONE + self.calculate_kv_scales = False + with ( set_forward_context( attn_metadata, self.vllm_config, - num_tokens=num_input_tokens, + num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, - cudagraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=None, + cudagraph_runtime_mode=cudagraph_mode, + batch_descriptor=batch_desc, ubatch_slices=ubatch_slices, ), record_function_or_nullcontext("Forward"), @@ -212,27 +134,32 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, - multimodal_kwargs=model_kwargs, + model_kwargs=model_kwargs, logits_indices=logits_indices, ) _, multimodal_outputs = self.extract_multimodal_outputs(outputs) - # Ensure one tensor per request, map to CPU for output struct - pooler_output: list[torch.Tensor | None] = [] + pooler_output: list[object] = [] if isinstance(multimodal_outputs, torch.Tensor): - # If model returned a single stacked tensor, split by requests + assert multimodal_outputs.shape[0] == 1, ( + "model should return a single tensor, to return multiple tensors, use a dict" + ) 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({"model_outputs": multimodal_outputs[i].detach().to("cpu").contiguous()}) elif isinstance(multimodal_outputs, list): + assert len(multimodal_outputs) == 1, ( + "model should return a single list, to return multiple lists, use a dict" + ) for out in multimodal_outputs: - pooler_output.append(out.detach().to("cpu").contiguous() if out is not None else None) + pooler_output.append( + {"model_outputs": 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) + for key, out in multimodal_outputs.items(): + pooler_output.append({key: out.detach().to("cpu").contiguous() if out is not None else None}) else: raise RuntimeError("Unsupported diffusion output type") - output = OmniModelRunnerOutput( req_ids=self.input_batch.req_ids, req_id_to_index=self.input_batch.req_id_to_index, @@ -257,11 +184,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): def _run_generation_model( self, *, - input_ids: torch.Tensor, - positions: torch.Tensor, + input_ids: torch.Tensor | None, + positions: torch.Tensor | None, intermediate_tensors: IntermediateTensors | None, inputs_embeds: torch.Tensor | None, - multimodal_kwargs: dict, + model_kwargs: dict, logits_indices: torch.Tensor, ) -> torch.Tensor | list[torch.Tensor]: """Run generation from codec codes to waveforms. @@ -279,7 +206,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, - **MultiModalKwargs.as_kwargs(multimodal_kwargs, device=self.device), + **MultiModalKwargs.as_kwargs(model_kwargs, device=self.device), sampling_metadata=self.input_batch.sampling_metadata, logits_index=logits_indices, sampler=self.sampler, @@ -305,10 +232,13 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): cudagraph_runtime_mode: CUDAGraphMode | None = None, force_attention: bool = False, uniform_decode: bool = False, + allow_microbatching: bool = True, skip_eplb: bool = False, is_profile: bool = False, create_mixed_batch: bool = False, remove_lora: bool = True, + activate_lora: bool = False, + is_graph_capturing: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """ Run a dummy forward pass to warm up/profile run or capture the @@ -331,12 +261,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): create_mixed_batch: If True, create a mixed batch with both decode (1 token) and prefill (multiple tokens) requests. remove_lora: If False, dummy LoRAs are not destroyed after the run + activate_lora: If False, dummy_run is performed without LoRAs. """ - assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in { - CUDAGraphMode.NONE, - CUDAGraphMode.PIECEWISE, - CUDAGraphMode.FULL, - } + assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes() # If cudagraph_mode.decode_mode() == FULL and # cudagraph_mode.separate_routine(). This means that we are using @@ -358,37 +285,79 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # has num_tokens in total. assert num_tokens <= self.scheduler_config.max_num_batched_tokens max_num_reqs = self.scheduler_config.max_num_seqs + if create_mixed_batch: + assert not uniform_decode + # Create mixed batch: + # first half decode tokens, second half one prefill + num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2) + num_prefill_tokens = num_tokens - num_decode_tokens + num_reqs = num_decode_tokens + 1 - num_reqs = min(num_tokens, max_num_reqs) - min_tokens_per_req = num_tokens // num_reqs - num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs - num_scheduled_tokens_list[-1] += num_tokens % num_reqs + # Create decode requests (1 token each) followed by prefill request + num_scheduled_tokens_list = [1] * num_decode_tokens + [num_prefill_tokens] + # Note: Overriding max_query_len to be the prefill tokens + max_query_len = num_prefill_tokens + elif uniform_decode: + assert not create_mixed_batch + num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len)) + num_scheduled_tokens_list = [max_query_len] * num_reqs + if num_tokens % max_query_len != 0: + num_scheduled_tokens_list[-1] = num_tokens % max_query_len + else: + num_reqs = min(num_tokens, max_num_reqs) + min_tokens_per_req = num_tokens // num_reqs + num_scheduled_tokens_list = [min_tokens_per_req] * num_reqs + num_scheduled_tokens_list[-1] += num_tokens % num_reqs assert sum(num_scheduled_tokens_list) == num_tokens assert len(num_scheduled_tokens_list) == num_reqs num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) + num_tokens_unpadded = int(num_scheduled_tokens.sum()) - num_tokens_after_padding = None + num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) - # If we failed to microbatch, currently need to resynchronize - # TODO(lucas,sage): we should be able to avoid this second sync by - # refactoring `get_dp_padding_ubatch` and `get_dp_padding` into - # a single `coordinate_batch_across_dp` function. - if num_tokens_after_padding is None: - num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) - num_tokens_after_padding = num_tokens + num_pad + _cudagraph_mode, batch_desc, ubatch_slices, num_tokens_across_dp = self._determine_batch_execution_and_padding( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs, + num_scheduled_tokens_np=num_scheduled_tokens, + max_num_scheduled_tokens=max_query_len, + use_cascade_attn=False, + allow_microbatching=allow_microbatching, + force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE), + # `force_uniform_decode` is used for cudagraph capture; because for + # capturing mixed prefill-decode batches, we sometimes use + # num_tokens == num_reqs which looks like a uniform decode batch to the + # dispatcher; but we actually want to capture a piecewise cudagraph + force_uniform_decode=uniform_decode, + # `force_has_lora` is used for cudagraph capture; because LoRA is + # activated later in the context manager, but we need to know the + # LoRA state when determining the batch descriptor for capture + force_has_lora=activate_lora, + ) + + if cudagraph_runtime_mode is None: + cudagraph_runtime_mode = _cudagraph_mode else: - num_tokens_across_dp = num_tokens_after_padding - num_tokens_after_padding = int(num_tokens_after_padding[0].item()) + assert cudagraph_runtime_mode == _cudagraph_mode, ( + f"Cudagraph runtime mode mismatch in dummy_run. " + f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}." + ) + + num_tokens_padded = batch_desc.num_tokens + num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs attn_metadata: PerLayerAttnMetadata | None = None # If force_attention is True, we always capture attention. Otherwise, # it only happens for cudagraph_runtime_mode=FULL. if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: - attn_metadata = {} - - seq_lens = max_query_len + if create_mixed_batch: + # In the mixed batch mode (used for FI warmup), we use + # shorter sequence lengths to run faster. + # TODO(luka) better system for describing dummy batches + seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] + else: + seq_lens = max_query_len # type: ignore[assignment] self.seq_lens.np[:num_reqs] = seq_lens self.seq_lens.np[num_reqs:] = 0 self.seq_lens.copy_to_gpu() @@ -397,50 +366,45 @@ class GPUGenerationModelRunner(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): - common_attn_metadata = CommonAttentionMetadata( - query_start_loc=self.query_start_loc.gpu[: num_reqs + 1], - query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1], - seq_lens=self.seq_lens.gpu[:num_reqs], - seq_lens_cpu=self.seq_lens.cpu[:num_reqs], - num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs], # noqa: E501 - num_reqs=num_reqs, - num_actual_tokens=num_tokens, - max_query_len=max_query_len, - max_seq_len=self.max_model_len, - block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs), - slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens], - causal=True, - ) - for attn_group in self.attn_groups[kv_cache_group_id]: - assert type(attn_metadata) is dict - attn_metadata_i = attn_group.get_metadata_builder().build_for_cudagraph_capture( - common_attn_metadata - ) # noqa: E501 - for layer_name in attn_group.layer_names: - attn_metadata[layer_name] = attn_metadata_i + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=ubatch_slices, + for_cudagraph_capture=is_graph_capturing, + ) - with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora): - model_kwargs = self._init_model_kwargs(num_tokens) + with self.maybe_dummy_run_with_lora( + self.lora_config, + num_scheduled_tokens, + num_sampled_tokens, + activate_lora, + remove_lora, + ): + # Make sure padding doesn't exceed max_num_tokens + assert num_tokens_padded <= self.max_num_tokens + model_kwargs = self._init_model_kwargs(num_tokens_padded) if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: input_ids = None - inputs_embeds = self.inputs_embeds.gpu[:num_tokens] + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] model_kwargs = { **model_kwargs, **self._dummy_mm_kwargs(num_reqs), } elif self.enable_prompt_embeds: input_ids = None - inputs_embeds = self.inputs_embeds.gpu[:num_tokens] - model_kwargs = self._init_model_kwargs(num_tokens) + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + model_kwargs = self._init_model_kwargs(num_tokens_padded) else: - input_ids = self.input_ids.gpu[:num_tokens] + input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None if self.uses_mrope: - positions = self.mrope_positions.gpu[:, :num_tokens] + positions = self.mrope_positions.gpu[:, :num_tokens_padded] + elif self.uses_xdrope_dim > 0: + positions = self.xdrope_positions.gpu[:, :num_tokens_padded] else: - positions = self.positions.gpu[:num_tokens] + positions = self.positions.gpu[:num_tokens_padded] if get_pp_group().is_first_rank: intermediate_tensors = None @@ -452,38 +416,26 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): 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_padded, None, False) - # filter out the valid batch descriptor - _cg_mode, batch_descriptor = ( - self.cudagraph_dispatcher.dispatch( - BatchDescriptor( - num_tokens=num_tokens_after_padding, - uniform_decode=uniform_decode, - ) - ) - if not is_profile - else (CUDAGraphMode.NONE, None) - ) - if cudagraph_runtime_mode is not None: - # we allow forcing NONE when the dispatcher disagrees to support - # warm ups for cudagraph capture - assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, ( - f"Cudagraph runtime mode mismatch at dummy_run. " - f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}." - ) - else: - cudagraph_runtime_mode = _cg_mode + if ubatch_slices is not None: + # Adjust values to reflect a single ubatch. + # TODO(sage,lucas): this is cruft that should be addressed in + # the padding refactor. + num_tokens_padded = ubatch_slices[0].num_tokens + if num_tokens_across_dp is not None: + num_tokens_across_dp[:] = num_tokens_padded with ( self.maybe_randomize_inputs(input_ids), set_forward_context( attn_metadata, self.vllm_config, - num_tokens=num_tokens_after_padding, + num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, + batch_descriptor=batch_desc, + ubatch_slices=ubatch_slices, ), ): outputs = self.model( @@ -498,10 +450,26 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): hidden_states, _ = outputs else: hidden_states = outputs - + hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) if self.speculative_config and self.speculative_config.use_eagle(): assert isinstance(self.drafter, EagleProposer) - self.drafter.dummy_run(num_tokens) + use_cudagraphs = ( + cudagraph_runtime_mode.has_mode(CUDAGraphMode.PIECEWISE) + and not self.speculative_config.enforce_eager + ) + + # Note(gnovack) - We need to disable cudagraphs for one of the two + # lora cases when cudagraph_specialize_lora is enabled. This is a + # short term mitigation for issue mentioned in + # https://github.com/vllm-project/vllm/issues/28334 + if self.compilation_config.cudagraph_specialize_lora and activate_lora: + use_cudagraphs = False + + self.drafter.dummy_run( + num_tokens, + use_cudagraphs=use_cudagraphs, + is_graph_capturing=is_graph_capturing, + ) # This is necessary to avoid blocking DP. # For dummy runs, we typically skip EPLB since we don't have any real @@ -510,15 +478,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # not have any requests to process, so they're executing dummy batches. # In such cases, we still have to trigger EPLB to make sure # ranks execute the rearrangement in synchronization. + if not skip_eplb: + self.eplb_step(is_dummy=True, is_profile=is_profile) - # logit_indices = np.cumsum(num_scheduled_tokens) - 1 - hidden_states, _ = self.extract_multimodal_outputs(hidden_states) return hidden_states, None def profile_run(self) -> None: # Profile with multimodal encoder & encoder cache. if self.supports_mm_inputs: - if self.model_config.multimodal_config.skip_mm_profiling: + mm_config = self.model_config.multimodal_config + if mm_config is not None and mm_config.skip_mm_profiling: logger.info("Skipping memory profiling for multimodal encoder and encoder cache.") else: mm_budget = self.mm_budget @@ -532,8 +501,9 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): 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, @@ -546,7 +516,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): ) # Run multimodal encoder. - dummy_encoder_outputs = self.model.get_multimodal_embeddings(**batched_dummy_mm_inputs) + dummy_encoder_outputs = self.model.embed_multimodal(**batched_dummy_mm_inputs) sanity_check_mm_encoder_outputs( dummy_encoder_outputs, @@ -556,13 +526,15 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # NOTE: This happens when encoder cache needs to store # the embeddings that encoder outputs are scattered onto. # In this case we create dummy embeddings of size - # (encode_budget, hidden_size) and scatter encoder - # output into it. + # (max_tokens_for_modality, hidden_size) and scatter + # encoder output into it. encoder_output_shape = dummy_encoder_outputs[0].shape - if encoder_output_shape[0] < encoder_budget: + max_mm_tokens_per_item = mm_budget.max_tokens_by_modality[dummy_modality] + if encoder_output_shape[0] < max_mm_tokens_per_item: + encoder_hidden_size = encoder_output_shape[-1] expanded_outputs = [] for output in dummy_encoder_outputs: - expanded = output.new_zeros((encoder_budget, encoder_output_shape[-1])) + expanded = output.new_zeros((max_mm_tokens_per_item, encoder_hidden_size)) num_tokens = output.shape[0] expanded[:num_tokens].copy_(output) expanded_outputs.append(expanded) @@ -574,8 +546,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): # Add `is_profile` here to pre-allocate communication buffers hidden_states, _ = self._dummy_run(self.max_num_tokens, is_profile=True) - if get_pp_group().is_last_rank: - pass + output = None self._sync_device() del hidden_states self.encoder_cache.clear() diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 49dc10c7b..27111f394 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -4,7 +4,8 @@ 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.utils.mem_constants import GiB_bytes +from vllm.utils.mem_utils import 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 @@ -20,18 +21,37 @@ class GPUGenerationWorker(GPUWorker): """ def init_device(self): - """Initialize CUDA device and distributed environment.""" - if self.device_config.device.type == "cuda": - # torch.distributed.all_reduce does not free the input tensor until - # the synchronization point. This causes the memory usage to grow - # as the number of all_reduce calls increases. This env var disables - # this behavior. - # Related issue: - # https://discuss.pytorch.org/t/cuda-allocation-lifetime-for-inputs-to-distributed-all-reduce/191573 - os.environ["TORCH_NCCL_AVOID_RECORD_STREAMS"] = "1" - + device = self.device_config.device + if isinstance(device, torch.device) and device.type == "cuda": # This env var set by Ray causes exceptions with graph building. os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + if ( + self.parallel_config.data_parallel_size > 1 + and self.parallel_config.data_parallel_size_local > 0 + and self.parallel_config.distributed_executor_backend not in ["ray", "external_launcher"] + and self.vllm_config.parallel_config.data_parallel_backend != "ray" + and self.vllm_config.parallel_config.nnodes_within_dp == 1 + ): + # Use local DP rank if available, otherwise use global DP rank. + dp_local_rank = self.parallel_config.data_parallel_rank_local + if dp_local_rank is None: + dp_local_rank = self.parallel_config.data_parallel_rank + + tp_pp_world_size = ( + self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size + ) + + # DP_LOCAL_RANK * TP_PP_WORLD_SIZE + TP_LOCAL_RANK + self.local_rank += dp_local_rank * tp_pp_world_size + assert self.local_rank < torch.cuda.device_count(), ( + f"DP adjusted local rank {self.local_rank} is out of bounds. " + ) + visible_device_count = torch.cuda.device_count() if torch.cuda.is_available() else 0 + assert self.parallel_config.local_world_size <= visible_device_count, ( + f"local_world_size ({self.parallel_config.local_world_size}) must " + f"be less than or equal to the number of visible devices " + f"({visible_device_count})." + ) self.device = torch.device(f"cuda:{self.local_rank}") current_platform.set_device(self.device) @@ -61,23 +81,23 @@ class GPUGenerationWorker(GPUWorker): self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization if self.init_snapshot.free_memory < self.requested_memory: - def GiB(b): - return round(b / GiB_bytes, 2) + def gib(bytes_val: float) -> float: + return round(bytes_val / GiB_bytes, 2) raise ValueError( f"Free memory on device " - f"({GiB(self.init_snapshot.free_memory)}/" - f"{GiB(self.init_snapshot.total_memory)} GiB) on startup " + f"({gib(self.init_snapshot.free_memory)}/" + f"{gib(self.init_snapshot.total_memory)} GiB) on startup " f"is less than desired GPU memory utilization " f"({self.cache_config.gpu_memory_utilization}, " - f"{GiB(self.requested_memory)} GiB). Decrease GPU memory " + f"{gib(self.requested_memory)} GiB). Decrease GPU memory " f"utilization or reduce GPU memory used by other processes." ) else: raise RuntimeError(f"Not support device type: {self.device_config.device}") # Construct the model runner - self.model_runner: GPUGenerationModelRunner = GPUGenerationModelRunner(self.vllm_config, self.device) + self.model_runner = GPUGenerationModelRunner(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 74d2e32f6..60dd271ab 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -4,22 +4,17 @@ 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 +from vllm.forward_context import set_forward_context from vllm.logger import init_logger from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding from vllm.model_executor.models.interfaces import supports_mrope from vllm.model_executor.models.interfaces_base import VllmModelForPooling from vllm.sampling_params import SamplingType -from vllm.utils import LazyLoader, cdiv -from vllm.v1.attention.backends.utils import ( - CommonAttentionMetadata, - split_attn_metadata, -) +from vllm.utils.import_utils import LazyLoader +from vllm.utils.math_utils import cdiv 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, IntermediateTensors, PerLayerAttnMetadata -from vllm.v1.worker.ubatch_splitting import ubatch_split -from vllm.v1.worker.ubatch_utils import UBatchSlices if TYPE_CHECKING: from vllm.v1.core.sched.output import SchedulerOutput @@ -35,6 +30,12 @@ logger = init_logger(__name__) class OmniGPUModelRunner(GPUModelRunner): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self._omni_per_req_additional_information: dict[str, dict] | None = None + self._omni_num_scheduled_tokens_np: np.ndarray | None = None + self._omni_last_model_output: object | None = None + def _init_mrope_positions(self, req_state: CachedRequestState): """Initialize M-RoPE positions for multimodal inputs. @@ -334,6 +335,8 @@ class OmniGPUModelRunner(GPUModelRunner): is_profile: bool = False, create_mixed_batch: bool = False, remove_lora: bool = True, + activate_lora: bool = False, + is_graph_capturing: bool = False, ) -> tuple[torch.Tensor, torch.Tensor]: """ Run a dummy forward pass to warm up/profile run or capture the @@ -356,12 +359,9 @@ class OmniGPUModelRunner(GPUModelRunner): create_mixed_batch: If True, create a mixed batch with both decode (1 token) and prefill (multiple tokens) requests. remove_lora: If False, dummy LoRAs are not destroyed after the run + activate_lora: If False, dummy_run is performed without LoRAs. """ - assert cudagraph_runtime_mode is None or cudagraph_runtime_mode in { - CUDAGraphMode.NONE, - CUDAGraphMode.PIECEWISE, - CUDAGraphMode.FULL, - } + assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes() # If cudagraph_mode.decode_mode() == FULL and # cudagraph_mode.separate_routine(). This means that we are using @@ -387,7 +387,7 @@ class OmniGPUModelRunner(GPUModelRunner): assert not uniform_decode # Create mixed batch: # first half decode tokens, second half one prefill - num_decode_tokens = num_tokens // 2 + num_decode_tokens = min(max_num_reqs - 1, num_tokens // 2) num_prefill_tokens = num_tokens - num_decode_tokens num_reqs = num_decode_tokens + 1 @@ -397,7 +397,7 @@ class OmniGPUModelRunner(GPUModelRunner): max_query_len = num_prefill_tokens elif uniform_decode: assert not create_mixed_batch - num_reqs = cdiv(num_tokens, max_query_len) + num_reqs = min(max_num_reqs, cdiv(num_tokens, max_query_len)) num_scheduled_tokens_list = [max_query_len] * num_reqs if num_tokens % max_query_len != 0: num_scheduled_tokens_list[-1] = num_tokens % max_query_len @@ -410,56 +410,52 @@ class OmniGPUModelRunner(GPUModelRunner): assert sum(num_scheduled_tokens_list) == num_tokens assert len(num_scheduled_tokens_list) == num_reqs num_scheduled_tokens = np.array(num_scheduled_tokens_list, dtype=np.int32) + num_tokens_unpadded = int(num_scheduled_tokens.sum()) - total_num_scheduled_tokens = int(num_scheduled_tokens.sum()) + num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) - ubatch_slices = None - num_tokens_after_padding = None + _cudagraph_mode, batch_desc, ubatch_slices, num_tokens_across_dp = self._determine_batch_execution_and_padding( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs, + num_scheduled_tokens_np=num_scheduled_tokens, + max_num_scheduled_tokens=max_query_len, + use_cascade_attn=False, + allow_microbatching=allow_microbatching, + force_eager=is_profile or (cudagraph_runtime_mode == CUDAGraphMode.NONE), + # `force_uniform_decode` is used for cudagraph capture; because for + # capturing mixed prefill-decode batches, we sometimes use + # num_tokens == num_reqs which looks like a uniform decode batch to the + # dispatcher; but we actually want to capture a piecewise cudagraph + force_uniform_decode=uniform_decode, + # `force_has_lora` is used for cudagraph capture; because LoRA is + # activated later in the context manager, but we need to know the + # LoRA state when determining the batch descriptor for capture + force_has_lora=activate_lora, + ) - # We currently only microbatch if the number of tokens is - # over a certain threshold. - if self.parallel_config.enable_dbo and allow_microbatching: - ubatch_slices, ubatch_num_tokens_after_padding = ubatch_split( - num_scheduled_tokens, - total_num_scheduled_tokens, - total_num_scheduled_tokens, - uniform_decode=uniform_decode, - vllm_config=self.vllm_config, - ) - # Currently when DBO is enabled `ubatch_split` returns - # the num_tokens_after_padding for a single ubatch, but we have 2 - # TODO(sage,lucas): this is cruft that should be addressed in the - # padding refactor. - if ubatch_num_tokens_after_padding is not None: - num_tokens_after_padding = ubatch_num_tokens_after_padding * 2 - - # If we failed to microbatch, currently need to resynchronize - # TODO(lucas,sage): we should be able to avoid this second sync by - # refactoring `get_dp_padding_ubatch` and `get_dp_padding` into - # a single `coordinate_batch_across_dp` function. - if num_tokens_after_padding is None: - num_pad, num_tokens_across_dp = self.get_dp_padding(num_tokens) - num_tokens_after_padding = num_tokens + num_pad + if cudagraph_runtime_mode is None: + cudagraph_runtime_mode = _cudagraph_mode else: - num_tokens_across_dp = num_tokens_after_padding - num_tokens_after_padding = int(num_tokens_after_padding[0].item()) + assert cudagraph_runtime_mode == _cudagraph_mode, ( + f"Cudagraph runtime mode mismatch in dummy_run. " + f"Expected {_cudagraph_mode}, but got {cudagraph_runtime_mode}." + ) + + num_tokens_padded = batch_desc.num_tokens + num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs attn_metadata: PerLayerAttnMetadata | None = None # If force_attention is True, we always capture attention. Otherwise, # it only happens for cudagraph_runtime_mode=FULL. if force_attention or cudagraph_runtime_mode == CUDAGraphMode.FULL: - attn_metadata = {} - if ubatch_slices is not None: - attn_metadata = [dict() for _ in range(len(ubatch_slices))] - if create_mixed_batch: # In the mixed batch mode (used for FI warmup), we use # shorter sequence lengths to run faster. # TODO(luka) better system for describing dummy batches seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] else: - seq_lens = max_query_len + seq_lens = max_query_len # type: ignore[assignment] self.seq_lens.np[:num_reqs] = seq_lens self.seq_lens.np[num_reqs:] = 0 self.seq_lens.copy_to_gpu() @@ -468,61 +464,45 @@ 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): - common_attn_metadata = CommonAttentionMetadata( - query_start_loc=self.query_start_loc.gpu[: num_reqs + 1], - query_start_loc_cpu=self.query_start_loc.cpu[: num_reqs + 1], - seq_lens=self.seq_lens.gpu[:num_reqs], - seq_lens_cpu=self.seq_lens.cpu[:num_reqs], - num_computed_tokens_cpu=self.input_batch.num_computed_tokens_cpu_tensor[:num_reqs], # noqa: E501 - num_reqs=num_reqs, - num_actual_tokens=num_tokens, - max_query_len=max_query_len, - max_seq_len=self.max_model_len, - block_table_tensor=self.input_batch.block_table[kv_cache_group_id].get_device_tensor(num_reqs), - slot_mapping=self.input_batch.block_table[kv_cache_group_id].slot_mapping.gpu[:num_tokens], - causal=True, - ) - for attn_group in self.attn_groups[kv_cache_group_id]: - 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): - assert common_attn_metadata.max_query_len == 1 - attn_metadata_i = attn_group.get_metadata_builder( - ubatch_id=ubid - ).build_for_cudagraph_capture(common_attn_metadata) - for layer_name in attn_group.layer_names: - assert type(attn_metadata) is list - 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( - common_attn_metadata - ) # noqa: E501 - for layer_name in attn_group.layer_names: - attn_metadata[layer_name] = attn_metadata_i + attn_metadata, _ = self._build_attention_metadata( + num_tokens=num_tokens_unpadded, + num_reqs=num_reqs_padded, + max_query_len=max_query_len, + ubatch_slices=ubatch_slices, + for_cudagraph_capture=is_graph_capturing, + ) - with self.maybe_dummy_run_with_lora(self.lora_config, num_scheduled_tokens, remove_lora): - model_kwargs = self._init_model_kwargs(num_tokens) + with self.maybe_dummy_run_with_lora( + self.lora_config, + num_scheduled_tokens, + num_sampled_tokens, + activate_lora, + remove_lora, + ): + # Make sure padding doesn't exceed max_num_tokens + assert num_tokens_padded <= self.max_num_tokens + model_kwargs = self._init_model_kwargs(num_tokens_padded) if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: input_ids = None - inputs_embeds = self.inputs_embeds.gpu[:num_tokens] + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] model_kwargs = { **model_kwargs, **self._dummy_mm_kwargs(num_reqs), } elif self.enable_prompt_embeds: input_ids = None - inputs_embeds = self.inputs_embeds.gpu[:num_tokens] - model_kwargs = self._init_model_kwargs(num_tokens) + inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + model_kwargs = self._init_model_kwargs(num_tokens_padded) else: - input_ids = self.input_ids.gpu[:num_tokens] + input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None if self.uses_mrope: - positions = self.mrope_positions.gpu[:, :num_tokens] + positions = self.mrope_positions.gpu[:, :num_tokens_padded] + elif self.uses_xdrope_dim > 0: + positions = self.xdrope_positions.gpu[:, :num_tokens_padded] else: - positions = self.positions.gpu[:num_tokens] + positions = self.positions.gpu[:num_tokens_padded] if get_pp_group().is_first_rank: intermediate_tensors = None @@ -534,46 +514,25 @@ class OmniGPUModelRunner(GPUModelRunner): device=self.device, ) - intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens, None, False) - - # filter out the valid batch descriptor - _cg_mode, batch_descriptor = ( - self.cudagraph_dispatcher.dispatch( - BatchDescriptor( - num_tokens=num_tokens_after_padding, - uniform_decode=uniform_decode, - ) - ) - if not is_profile - else (CUDAGraphMode.NONE, None) - ) - if cudagraph_runtime_mode is not None: - # we allow forcing NONE when the dispatcher disagrees to support - # warm ups for cudagraph capture - assert cudagraph_runtime_mode == CUDAGraphMode.NONE or cudagraph_runtime_mode == _cg_mode, ( - f"Cudagraph runtime mode mismatch at dummy_run. " - f"Expected {_cg_mode}, but got {cudagraph_runtime_mode}." - ) - else: - cudagraph_runtime_mode = _cg_mode + intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False) if ubatch_slices is not None: # Adjust values to reflect a single ubatch. # TODO(sage,lucas): this is cruft that should be addressed in # the padding refactor. - num_tokens_after_padding = ubatch_slices[0].num_tokens + num_tokens_padded = ubatch_slices[0].num_tokens if num_tokens_across_dp is not None: - num_tokens_across_dp[:] = num_tokens_after_padding + num_tokens_across_dp[:] = num_tokens_padded with ( self.maybe_randomize_inputs(input_ids), set_forward_context( attn_metadata, self.vllm_config, - num_tokens=num_tokens_after_padding, + num_tokens=num_tokens_padded, num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=cudagraph_runtime_mode, - batch_descriptor=batch_descriptor, + batch_descriptor=batch_desc, ubatch_slices=ubatch_slices, ), ): @@ -589,10 +548,26 @@ class OmniGPUModelRunner(GPUModelRunner): hidden_states, _ = outputs else: hidden_states = outputs - + hidden_states, multimodal_outputs = self.extract_multimodal_outputs(hidden_states) if self.speculative_config and self.speculative_config.use_eagle(): assert isinstance(self.drafter, EagleProposer) - self.drafter.dummy_run(num_tokens) + use_cudagraphs = ( + cudagraph_runtime_mode.has_mode(CUDAGraphMode.PIECEWISE) + and not self.speculative_config.enforce_eager + ) + + # Note(gnovack) - We need to disable cudagraphs for one of the two + # lora cases when cudagraph_specialize_lora is enabled. This is a + # short term mitigation for issue mentioned in + # https://github.com/vllm-project/vllm/issues/28334 + if self.compilation_config.cudagraph_specialize_lora and activate_lora: + use_cudagraphs = False + + self.drafter.dummy_run( + num_tokens, + use_cudagraphs=use_cudagraphs, + is_graph_capturing=is_graph_capturing, + ) # This is necessary to avoid blocking DP. # For dummy runs, we typically skip EPLB since we don't have any real @@ -605,110 +580,8 @@ 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) - return hidden_states, hidden_states[logit_indices] - - def _preprocess( - self, - scheduler_output: "SchedulerOutput", - num_scheduled_tokens_np: np.ndarray, - intermediate_tensors: IntermediateTensors | None = None, - ubatch_slices: UBatchSlices | None = None, - num_tokens_after_padding: torch.Tensor | None = None, - ) -> tuple[ - int, - int, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor | None, - torch.Tensor, - IntermediateTensors | None, - dict[str, Any], - dict[str, dict] | None, - ]: - num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens - if ubatch_slices: - assert num_tokens_after_padding is not None - num_input_tokens = int(num_tokens_after_padding[0].item() * 2) - self.pad_out_ubatch_slice(ubatch_slices, num_input_tokens) - elif ubatch_slices is None: - num_input_tokens = self._get_num_input_tokens(num_scheduled_tokens) - num_pad, num_tokens_after_padding = self.get_dp_padding(num_input_tokens) - num_input_tokens += num_pad - - # _prepare_inputs may reorder the batch, so we must gather multi - # modal outputs after that to ensure the correct order - per_req_additional_information: dict[str, dict] | None = None - if self.supports_mm_inputs and get_pp_group().is_first_rank and not self.model_config.is_encoder_decoder: - # Build multimodal inputs and overlay prompt embeds; collect per-request info - per_req_additional_information = self._build_mm_inputs_and_overlays( - scheduler_output, num_scheduled_tokens, num_scheduled_tokens_np - ) - - input_ids = self.input_ids.gpu[:num_input_tokens] - inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] - model_kwargs = { - **self._init_model_kwargs(num_scheduled_tokens), - **self._extract_mm_kwargs(scheduler_output), - } - elif self.enable_prompt_embeds and get_pp_group().is_first_rank: - # Get the input embeddings for the tokens that are not input embeds, - # then put them into the appropriate positions. - # TODO(qthequartermasterman): Since even when prompt embeds are - # enabled, (a) not all requests will use prompt embeds, and (b) - # after the initial prompt is processed, the rest of the generated - # tokens will be token ids, it is not desirable to have the - # embedding layer outside of the CUDA graph all the time. The v0 - # engine avoids this by "double compiling" the CUDA graph, once - # with input_ids and again with inputs_embeds, for all num_tokens. - # If a batch only has token ids, then including the embedding layer - # in the CUDA graph will be more performant (like in the else case - # below). - token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1) - # Some tokens ids may need to become embeds - if token_ids_idx.numel() > 0: - token_ids = self.input_ids.gpu[token_ids_idx] - tokens_to_embeds = self.model.get_input_embeddings(input_ids=token_ids) - self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds - - inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] - model_kwargs = self._init_model_kwargs(num_input_tokens) - input_ids = None - else: - # For text-only models, we use token ids as input. - # While it is possible to use embeddings as input just like the - # multimodal models, it is not desirable for performance since - # then the embedding layer is not included in the CUDA graph. - input_ids = self.input_ids.gpu[:num_input_tokens] - inputs_embeds = None - model_kwargs = self._init_model_kwargs(num_input_tokens) - if self.uses_mrope: - positions = self.mrope_positions.gpu[:, :num_input_tokens] - else: - positions = self.positions.gpu[:num_input_tokens] - - if get_pp_group().is_first_rank: - intermediate_tensors = None - else: - intermediate_tensors = self.sync_and_slice_intermediate_tensors( - num_input_tokens, intermediate_tensors, True - ) - - if self.model_config.is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: - encoder_inputs = self._extract_encoder_inputs(scheduler_output) - model_kwargs.update(encoder_inputs) - - return ( - num_scheduled_tokens, - num_input_tokens, - num_tokens_after_padding, - input_ids, - inputs_embeds, - positions, - intermediate_tensors, - model_kwargs, - per_req_additional_information, - ) + logit_indices_device = torch.from_numpy(logit_indices).to(self.device, non_blocking=True) + return hidden_states, hidden_states[logit_indices_device] def _decode_and_store_request_payloads(self, scheduler_output: "SchedulerOutput") -> None: """Decode per-request prompt_embeds and additional_information for newly @@ -770,7 +643,14 @@ class OmniGPUModelRunner(GPUModelRunner): for req_id in self.input_batch.req_ids: req_state = self.requests.get(req_id) info = getattr(req_state, "additional_information_cpu", None) if req_state is not None else None - per_req_runtime_info.append(info if isinstance(info, dict) else {}) + if info and isinstance(info, dict): + per_req_runtime_info.append(info) + if "thinker_reply_part_per_request" in info: + q = info["thinker_reply_part_per_request"] + if hasattr(q, "shape"): + logger.debug(f"[OMNI] req={req_id} has thinker_reply_part_per_request queue shape: {q.shape}") + else: + per_req_runtime_info.append({}) return per_req_runtime_info def _compute_request_token_spans(self, num_scheduled_tokens_np) -> list[tuple[int, int]]: @@ -798,24 +678,34 @@ class OmniGPUModelRunner(GPUModelRunner): model_kwargs_extra["additional_information_by_req_id"] = per_req_additional_information try: model_kwargs_extra["runtime_additional_information"] = self._gather_runtime_additional_information() - model_kwargs_extra["request_ids"] = self.input_batch.req_ids - model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(num_scheduled_tokens_np) - except Exception: - pass + model_kwargs_extra["request_ids"] = list(self.input_batch.req_ids) + if num_scheduled_tokens_np is not None: + model_kwargs_extra["request_token_spans"] = self._compute_request_token_spans(num_scheduled_tokens_np) + except Exception as e: + logger.error(f"[OMNI DEBUG] Error building model_kwargs_extra: {e}") + import traceback + + traceback.print_exc() return model_kwargs_extra def _process_additional_information_updates(self, multimodal_outputs: object) -> None: """Process model-provided per-request additional_information updates and merge into request state.""" try: + if multimodal_outputs is None: + logger.debug("[OMNI] multimodal_outputs is None, no updates to process") + return if not isinstance(multimodal_outputs, dict): + logger.debug(f"[OMNI] multimodal_outputs is not dict: {type(multimodal_outputs)}") return if ( "additional_information_update" not in multimodal_outputs and "additional_information_update_by_req_id" not in multimodal_outputs ): + logger.debug(f"[OMNI] multimodal_outputs has no update keys: {list(multimodal_outputs.keys())}") return updates_list = multimodal_outputs.get("additional_information_update") if isinstance(updates_list, list): + logger.debug(f"[OMNI] Processing additional_information_update list with {len(updates_list)} items") for idx, upd in enumerate(updates_list): if not isinstance(upd, dict) or idx >= len(self.input_batch.req_ids): continue @@ -823,18 +713,33 @@ class OmniGPUModelRunner(GPUModelRunner): self._merge_additional_information_update(req_id, upd) updates_map = multimodal_outputs.get("additional_information_update_by_req_id") if isinstance(updates_map, dict): + logger.debug( + f"[OMNI] Processing additional_information_update_by_req_id for {list(updates_map.keys())}" + ) for req_id, upd in updates_map.items(): if not isinstance(upd, dict): continue if req_id not in self.requests: + logger.warning(f"[OMNI] req_id {req_id} not found in requests") continue self._merge_additional_information_update(req_id, upd) + + req_state = self.requests.get(req_id) + if req_state: + info = getattr(req_state, "additional_information_cpu", {}) + if "thinker_reply_part_per_request" in info: + q = info["thinker_reply_part_per_request"] + if hasattr(q, "shape"): + logger.debug(f"[OMNI] After update, req={req_id} queue shape: {q.shape}") except Exception as e: logger.error( f"Error merging for requests:{self.input_batch.req_ids} " f"additional information update: {e}, with the multimodal_outputs " f"as {multimodal_outputs}" ) + import traceback + + traceback.print_exc() def _build_mm_inputs_and_overlays( self, @@ -895,3 +800,164 @@ class OmniGPUModelRunner(GPUModelRunner): if req_info: per_req_additional_information[req_id] = req_info return per_req_additional_information + + def _preprocess( + self, + scheduler_output: "SchedulerOutput", + num_input_tokens: int, + intermediate_tensors: IntermediateTensors | None = None, + ): + """Align with v0.12 preprocess and omni's additional information handling.""" + # Decode payload first, ensure request state has prompt_embeds / additional_information + self._decode_and_store_request_payloads(scheduler_output) + + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + is_first_rank = get_pp_group().is_first_rank + is_encoder_decoder = self.model_config.is_encoder_decoder + + # _prepare_inputs may reorder the batch, so we must gather multi + # modal outputs after that to ensure the correct order + ec_connector_output = None + + if self.supports_mm_inputs and is_first_rank and not is_encoder_decoder: + # Run the multimodal encoder if any. + with self.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=self.encoder_cache, + ) as ec_connector_output: + self._execute_mm_encoder(scheduler_output) + mm_embeds, is_mm_embed = self._gather_mm_embeddings(scheduler_output) + + # NOTE(woosuk): To unify token ids and soft tokens (vision + # embeddings), we always use embeddings (rather than token ids) + # as input to the multimodal model, even when the input is text. + inputs_embeds_scheduled = self.model.embed_input_ids( + self.input_ids.gpu[:num_scheduled_tokens], + multimodal_embeddings=mm_embeds, + is_multimodal=is_mm_embed, + ) + + # TODO(woosuk): Avoid the copy. Optimize. + self.inputs_embeds.gpu[:num_scheduled_tokens].copy_(inputs_embeds_scheduled) + + input_ids = self.input_ids.gpu[:num_input_tokens] + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = { + **self._init_model_kwargs(num_scheduled_tokens), + **self._extract_mm_kwargs(scheduler_output), + } + elif self.enable_prompt_embeds and is_first_rank: + # Get the input embeddings for the tokens that are not input embeds, + # then put them into the appropriate positions. + # TODO(qthequartermasterman): Since even when prompt embeds are + # enabled, (a) not all requests will use prompt embeds, and (b) + # after the initial prompt is processed, the rest of the generated + # tokens will be token ids, it is not desirable to have the + # embedding layer outside of the CUDA graph all the time. The v0 + # engine avoids this by "double compiling" the CUDA graph, once + # with input_ids and again with inputs_embeds, for all num_tokens. + # If a batch only has token ids, then including the embedding layer + # in the CUDA graph will be more performant (like in the else case + # below). + token_ids_idx = self.is_token_ids.gpu[:num_scheduled_tokens].nonzero(as_tuple=False).squeeze(1) + # Some tokens ids may need to become embeds + if token_ids_idx.numel() > 0: + token_ids = self.input_ids.gpu[token_ids_idx] + tokens_to_embeds = self.model.embed_input_ids(input_ids=token_ids) + self.inputs_embeds.gpu[token_ids_idx] = tokens_to_embeds + + inputs_embeds = self.inputs_embeds.gpu[:num_input_tokens] + model_kwargs = self._init_model_kwargs(num_input_tokens) + input_ids = self.input_ids.gpu[:num_input_tokens] + else: + # For text-only models, we use token ids as input. + # While it is possible to use embeddings as input just like the + # multimodal models, it is not desirable for performance since + # then the embedding layer is not included in the CUDA graph. + input_ids = self.input_ids.gpu[:num_input_tokens] + inputs_embeds = None + model_kwargs = self._init_model_kwargs(num_input_tokens) + + if self.uses_mrope: + positions = self.mrope_positions.gpu[:, :num_input_tokens] + elif self.uses_xdrope_dim > 0: + positions = self.xdrope_positions.gpu[:, :num_input_tokens] + else: + positions = self.positions.gpu[:num_input_tokens] + + if is_first_rank: + intermediate_tensors = None + else: + assert intermediate_tensors is not None + intermediate_tensors = self.sync_and_slice_intermediate_tensors( + num_input_tokens, intermediate_tensors, True + ) + + if is_encoder_decoder and scheduler_output.scheduled_encoder_inputs: + # Run the encoder, just like we do with other multimodal inputs. + # For an encoder-decoder model, our processing here is a bit + # simpler, because the outputs are just passed to the decoder. + # We are not doing any prompt replacement. We also will only + # ever have a single encoder input. + encoder_outputs = self._execute_mm_encoder(scheduler_output) + model_kwargs.update({"encoder_outputs": encoder_outputs}) + + req_ids = self.input_batch.req_ids + num_scheduled_tokens_np = np.array( + [scheduler_output.num_scheduled_tokens[rid] for rid in req_ids], + dtype=np.int32, + ) + self._omni_num_scheduled_tokens_np = num_scheduled_tokens_np + + # Always collect additional_information, not only in prefill. + # _collect_additional_information_for_prefill handles prompt_embeds overlay during prefill, + # but even in decode we need per_req_additional_information so that + # _build_model_kwargs_extra can pass runtime_additional_information correctly. + per_req_additional_information: dict[str, dict] = {} + if inputs_embeds is not None: + # Prefill: overlay prompt_embeds and collect additional_information + per_req_additional_information = self._collect_additional_information_for_prefill(num_scheduled_tokens_np) + # Save for _model_forward regardless of prefill; runtime_additional_information + # will be fetched from request state inside _build_model_kwargs_extra + self._omni_per_req_additional_information = per_req_additional_information + + return ( + input_ids, + inputs_embeds, + positions, + intermediate_tensors, + model_kwargs, + ec_connector_output, + ) + + def _model_forward( + self, + input_ids: torch.Tensor | None = None, + positions: torch.Tensor | None = None, + intermediate_tensors: IntermediateTensors | None = None, + inputs_embeds: torch.Tensor | None = None, + **model_kwargs: dict[str, Any], + ): + """Inject omni-specific kwargs into forward and cache model output""" + model_kwargs_extra = self._build_model_kwargs_extra( + self._omni_per_req_additional_information, + self._omni_num_scheduled_tokens_np, + ) + + runtime_info = model_kwargs_extra.get("runtime_additional_information", []) + if runtime_info: + for i, info in enumerate(runtime_info): + if info: + logger.debug(f"[OMNI] req[{i}] runtime_additional_information keys: {list(info.keys())}") + + model_output = super()._model_forward( + input_ids=input_ids, + positions=positions, + intermediate_tensors=intermediate_tensors, + inputs_embeds=inputs_embeds, + **model_kwargs, + **model_kwargs_extra, + ) + # Cache model output so later sample_tokens can consume multimodal results. + self._omni_last_model_output = model_output + return model_output