From 776c3a794ecd0acb374745f6059fd4afa63dc39d Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Wed, 21 Jan 2026 00:25:25 +0800 Subject: [PATCH] Dev/rebase 0.14.0 (#847) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: JaredforReal Signed-off-by: root Signed-off-by: tzhouam Signed-off-by: Taichang Zhou Signed-off-by: TangPeng <85704592@qq.com> Signed-off-by: XU Mingshi <91017482+mxuax@users.noreply.github.com> Signed-off-by: mxuax Signed-off-by: Sihyeon Jang Signed-off-by: zjy0516 Signed-off-by: iwzbi Signed-off-by: ZeldaHuang Signed-off-by: Yuhan Liu <30294295+liuyuhanalex@users.noreply.github.com> Signed-off-by: Alicia <115451386+congw729@users.noreply.github.com> Signed-off-by: samithuang <285365963@qq.com> Signed-off-by: linyueqian Signed-off-by: David Chen <530634352@qq.com> Signed-off-by: yinpeiqi Signed-off-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Signed-off-by: Hongsheng Liu Signed-off-by: wangyu31577 Signed-off-by: princepride Signed-off-by: John Liu BUAA Signed-off-by: Dinesh G Signed-off-by: gDINESH13 Signed-off-by: root Signed-off-by: Zhou Taichang Co-authored-by: JaredforReal Co-authored-by: root Co-authored-by: JustQJ <37905360+JustQJ@users.noreply.github.com> Co-authored-by: XU Mingshi <91017482+mxuax@users.noreply.github.com> Co-authored-by: Sihyeon Jang Co-authored-by: Jiangyun Zhu Co-authored-by: catcat <108673086+iwzbi@users.noreply.github.com> Co-authored-by: Ziming Huang Co-authored-by: Yuhan Liu <30294295+liuyuhanalex@users.noreply.github.com> Co-authored-by: Alicia <115451386+congw729@users.noreply.github.com> Co-authored-by: Samit <285365963@qq.com> Co-authored-by: Yueqian Lin <70319226+linyueqian@users.noreply.github.com> Co-authored-by: WeiQing Chen <40507679+david6666666@users.noreply.github.com> Co-authored-by: Peiqi Yin <60515999+yinpeiqi@users.noreply.github.com> Co-authored-by: Didan Deng <33117903+wtomin@users.noreply.github.com> Co-authored-by: Hongsheng Liu Co-authored-by: wangyu <53896905+yenuo26@users.noreply.github.com> Co-authored-by: wangyu31577 Co-authored-by: 汪志鹏 Co-authored-by: John Liu BUAA Co-authored-by: D!NE$H <67671800+gDINESH13@users.noreply.github.com> Co-authored-by: root --- .buildkite/scripts/simple_test.sh | 3 +- README.md | 2 +- docker/Dockerfile.ci | 2 +- docker/Dockerfile.rocm | 2 +- docs/configuration/README.md | 2 +- docs/design/architecture_overview.md | 2 +- .../installation/gpu/cuda.inc.md | 6 +- .../installation/gpu/rocm.inc.md | 4 +- .../installation/npu/npu.inc.md | 8 +- .../offline_inference/qwen2_5_omni/end2end.py | 6 +- .../qwen2_5_omni/run_single_prompt.sh | 2 +- .../offline_inference/qwen3_omni/end2end.py | 73 +- ...letion_client_for_multimodal_generation.py | 27 + mkdocs.yml | 1 + pyproject.toml | 1 + tests/conftest.py | 32 +- .../attention/test_sequence_parallel.py | 5 + tests/diffusion/distributed/test_comm.py | 15 + .../stage_configs/qwen2_5_omni_ci.yaml | 1 + .../stage_configs/qwen3_omni_ci.yaml | 3 +- .../stage_configs/rocm/qwen3_omni_ci.yaml | 1 + .../offline_inference/test_qwen2_5_omni.py | 4 +- .../e2e/offline_inference/test_qwen3_omni.py | 4 +- .../test_sequence_parallel.py | 2 - .../test_stable_audio_model.py | 3 - tests/e2e/offline_inference/test_t2v_model.py | 8 +- .../test_zimage_tensor_parallel.py | 7 +- .../stage_configs/qwen3_omni_ci.yaml | 1 + .../stage_configs/rocm/qwen3_omni_ci.yaml | 1 + tests/e2e/online_serving/test_async_omni.py | 37 +- .../test_i2i_multi_image_input.py | 2 +- tests/e2e/online_serving/test_qwen3_omni.py | 9 +- .../test_qwen3_omni_expansion.py | 4 +- tests/e2e/stage_configs/qwen3_omni_ci.yaml | 1 + tests/worker/test_omni_gpu_model_runner.py | 2 +- vllm_omni/assets/video.py | 16 + vllm_omni/config/model.py | 142 +-- vllm_omni/core/sched/omni_ar_scheduler.py | 67 +- .../core/sched/omni_generation_scheduler.py | 8 +- vllm_omni/diffusion/forward_context.py | 11 +- .../diffusion/models/glm_image/__init__.py | 21 + .../models/glm_image/glm_image_transformer.py | 792 ++++++++++++++ .../models/glm_image/pipeline_glm_image.py | 981 ++++++++++++++++++ vllm_omni/diffusion/registry.py | 7 + vllm_omni/engine/__init__.py | 69 +- vllm_omni/engine/output_processor.py | 61 -- vllm_omni/entrypoints/async_omni_llm.py | 4 +- vllm_omni/entrypoints/chat_utils.py | 4 +- vllm_omni/entrypoints/omni_llm.py | 47 + vllm_omni/entrypoints/openai/api_server.py | 336 ++++-- .../openai/protocol/chat_completion.py | 4 +- vllm_omni/entrypoints/openai/serving_chat.py | 170 +-- .../models/qwen2_5_omni/qwen2_5_omni.py | 31 +- .../qwen2_5_omni/qwen2_5_omni_thinker.py | 47 +- .../models/qwen2_5_omni/qwen2_old.py | 73 +- .../models/qwen3_omni/qwen3_omni.py | 26 +- .../qwen3_omni/qwen3_omni_moe_thinker.py | 132 ++- .../stage_configs/qwen2_5_omni.yaml | 1 + .../stage_configs/qwen3_omni_moe.yaml | 7 +- vllm_omni/worker/gpu_ar_model_runner.py | 299 ++++-- vllm_omni/worker/gpu_ar_worker.py | 45 +- .../worker/gpu_generation_model_runner.py | 344 ++++-- vllm_omni/worker/gpu_generation_worker.py | 49 +- vllm_omni/worker/gpu_model_runner.py | 188 +++- 64 files changed, 3359 insertions(+), 906 deletions(-) create mode 100644 vllm_omni/assets/video.py create mode 100644 vllm_omni/diffusion/models/glm_image/__init__.py create mode 100644 vllm_omni/diffusion/models/glm_image/glm_image_transformer.py create mode 100644 vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py diff --git a/.buildkite/scripts/simple_test.sh b/.buildkite/scripts/simple_test.sh index 55ac27cec..576802bfe 100755 --- a/.buildkite/scripts/simple_test.sh +++ b/.buildkite/scripts/simple_test.sh @@ -47,8 +47,9 @@ 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.12.0 +"${UV_BIN}" pip install --python "${VENV_PYTHON}" vllm==0.14.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/ "${VENV_PYTHON}" -m pytest -v -s tests/model_executor/models/qwen2_5_omni/test_audio_length.py diff --git a/README.md b/README.md index fd2fed534..df377293c 100644 --- a/README.md +++ b/README.md @@ -15,7 +15,7 @@ Easy, fast, and cheap omni-modality model serving for everyone --- *Latest News* 🔥 - +- [2026/01] We released [0.14.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.14.0). - [2026/01] We released [0.12.0rc1](https://github.com/vllm-project/vllm-omni/releases/tag/v0.12.0rc1) - a major RC milestone focused on maturing the diffusion stack, strengthening OpenAI-compatible serving, expanding omni-model coverage, and improving stability across platforms (GPU/NPU/ROCm), please check our latest [design](https://docs.google.com/presentation/d/1qv4qMW1rKAqDREMXiUDLIgqqHQe7TDPj/edit?usp=sharing&ouid=110473603432222024453&rtpof=true&sd=true). - [2025/11] vLLM community officially released [vllm-project/vllm-omni](https://github.com/vllm-project/vllm-omni) in order to support omni-modality models serving. diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index aecc42945..f0f585fba 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.12.0 +ARG VLLM_BASE_TAG=v0.14.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 7fabb9c3c..5b480ff52 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -2,7 +2,7 @@ ARG BASE_IMAGE=rocm/vllm-dev:nightly_main_20251205 FROM ${BASE_IMAGE} ARG COMMON_WORKDIR=/app -ARG VLLM_VERSION=v0.12.0 +ARG VLLM_VERSION=v0.14.0 ARG PYTORCH_ROCM_ARCH="gfx942;gfx950" WORKDIR ${COMMON_WORKDIR} diff --git a/docs/configuration/README.md b/docs/configuration/README.md index 37e28bd0c..02440b95d 100644 --- a/docs/configuration/README.md +++ b/docs/configuration/README.md @@ -2,7 +2,7 @@ This section lists the most common options for running vLLM-Omni. -For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.12.0/configuration/index.html) +For options within a vLLM Engine. Please refer to [vLLM Configuration](https://docs.vllm.ai/en/v0.14.0/configuration/index.html) Currently, the main options are maintained by stage configs for each model. diff --git a/docs/design/architecture_overview.md b/docs/design/architecture_overview.md index 16f81ab15..6793895cd 100644 --- a/docs/design/architecture_overview.md +++ b/docs/design/architecture_overview.md @@ -192,4 +192,4 @@ curl -sS -X POST http://localhost:8091/v1/chat/completions \ } ``` -For more usages, please refer to [examples](../examples/README.md). +For more usages, please refer to [examples](https://github.com/vllm-project/vllm-omni/tree/main/examples). diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index 540b18529..8108bb6c8 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -58,11 +58,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.12.0 +git checkout v0.14.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.12.0/vllm-0.12.0-cp38-abi3-manylinux_2_31_x86_64.whl +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.14.0/vllm-0.14.0-cp38-abi3-manylinux_2_31_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash @@ -93,7 +93,7 @@ docker run --runtime nvidia --gpus 2 \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ --ipc=host \ - vllm/vllm-omni:v0.12.0rc1 \ + vllm/vllm-omni:v0.14.0 \ --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 1fa751e25..81b62487f 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -68,7 +68,7 @@ docker run -it \ -v :/app/model \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ - vllm/vllm-omni-rocm:v0.12.0rc1 \ + vllm/vllm-omni-rocm:v0.14.0 \ vllm serve --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` @@ -86,7 +86,7 @@ docker run -it \ -v :/app/model \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ - vllm/vllm-omni-rocm:v0.12.0rc1 \ + vllm/vllm-omni-rocm:v0.14.0 \ bash ``` diff --git a/docs/getting_started/installation/npu/npu.inc.md b/docs/getting_started/installation/npu/npu.inc.md index 9c36c2be6..9044bb089 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.12.0rc1 +# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0 # Atlas A3: -# export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1-a3 -export IMAGE=quay.io/ascend/vllm-ascend:v0.12.0rc1 +# export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0-a3 +export IMAGE=quay.io/ascend/vllm-ascend:v0.14.0 docker run --rm \ --name vllm-omni-npu \ --shm-size=1g \ @@ -42,7 +42,7 @@ source ~/.bashrc # Inside the container, install vLLM-Omni from source cd /vllm-workspace -git clone -b v0.12.0rc1 https://github.com/vllm-project/vllm-omni.git +git clone -b v0.14.0 https://github.com/vllm-project/vllm-omni.git cd vllm-omni pip install -v -e . export VLLM_WORKER_MULTIPROC_METHOD=spawn diff --git a/examples/offline_inference/qwen2_5_omni/end2end.py b/examples/offline_inference/qwen2_5_omni/end2end.py index d1e0264cd..22f2161c2 100644 --- a/examples/offline_inference/qwen2_5_omni/end2end.py +++ b/examples/offline_inference/qwen2_5_omni/end2end.py @@ -278,9 +278,9 @@ def get_audio_query(question: str = None, audio_path: str | None = None, samplin query_map = { - "mixed_modalities": get_mixed_modalities_query, + "use_mixed_modalities": get_mixed_modalities_query, "use_audio_in_video": get_use_audio_in_video_query, - "multi_audios": get_multi_audios_query, + "use_multi_audios": get_multi_audios_query, "use_image": get_image_query, "use_video": get_video_query, "use_audio": get_audio_query, @@ -434,7 +434,7 @@ def parse_args(): "--query-type", "-q", type=str, - default="mixed_modalities", + default="use_mixed_modalities", choices=query_map.keys(), help="Query type.", ) diff --git a/examples/offline_inference/qwen2_5_omni/run_single_prompt.sh b/examples/offline_inference/qwen2_5_omni/run_single_prompt.sh index 26cbf0172..c8e4cd2cb 100644 --- a/examples/offline_inference/qwen2_5_omni/run_single_prompt.sh +++ b/examples/offline_inference/qwen2_5_omni/run_single_prompt.sh @@ -1,2 +1,2 @@ python end2end.py --output-wav output_audio \ - --query-type mixed_modalities + --query-type use_mixed_modalities diff --git a/examples/offline_inference/qwen3_omni/end2end.py b/examples/offline_inference/qwen3_omni/end2end.py index 7280d7d2e..3cd891820 100644 --- a/examples/offline_inference/qwen3_omni/end2end.py +++ b/examples/offline_inference/qwen3_omni/end2end.py @@ -12,6 +12,7 @@ from typing import NamedTuple import librosa import numpy as np import soundfile as sf +import vllm from PIL import Image from vllm import SamplingParams from vllm.assets.audio import AudioAsset @@ -225,18 +226,76 @@ def get_multi_audios_query() -> QueryResult: ) +# def get_use_audio_in_video_query(video_path: str | None = None) -> QueryResult: +# question = ( +# "Describe the content of the video in details, then convert what the " +# "baby say into text." +# ) +# prompt = ( +# f"<|im_start|>system\n{default_system}<|im_end|>\n" +# "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>" +# f"{question}<|im_end|>\n" +# f"<|im_start|>assistant\n" +# ) +# if video_path: +# if not os.path.exists(video_path): +# raise FileNotFoundError(f"Video file not found: {video_path}") +# video_frames = video_to_ndarrays(video_path, num_frames=16) +# else: +# video_frames = VideoAsset(name="baby_reading", num_frames=16).np_ndarrays +# audio = extract_video_audio(video_path, sampling_rate=16000) +# return QueryResult( +# inputs={ +# "prompt": prompt, +# "multi_modal_data": { +# "video": video_frames, +# "audio": audio, +# }, +# "mm_processor_kwargs": { +# "use_audio_in_video": True, +# }, +# }, +# limit_mm_per_prompt={"audio": 1, "video": 1}, +# ) +def get_use_audio_in_video_query() -> QueryResult: + question = "Describe the content of the video in details, then convert what the baby say into text." + prompt = ( + f"<|im_start|>system\n{default_system}<|im_end|>\n" + "<|im_start|>user\n<|vision_start|><|video_pad|><|vision_end|>" + f"{question}<|im_end|>\n" + f"<|im_start|>assistant\n" + ) + asset = VideoAsset(name="baby_reading", num_frames=16) + audio = asset.get_audio(sampling_rate=16000) + return QueryResult( + inputs={ + "prompt": prompt, + "multi_modal_data": { + "video": asset.np_ndarrays, + "audio": audio, + }, + "mm_processor_kwargs": { + "use_audio_in_video": True, + }, + }, + limit_mm_per_prompt={"audio": 1, "video": 1}, + ) + + query_map = { "text": get_text_query, "use_audio": get_audio_query, "use_image": get_image_query, "use_video": get_video_query, - "multi_audios": get_multi_audios_query, - "mixed_modalities": get_mixed_modalities_query, + "use_multi_audios": get_multi_audios_query, + "use_mixed_modalities": get_mixed_modalities_query, + "use_audio_in_video": get_use_audio_in_video_query, } def main(args): model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct" + print("=" * 20, "\n", f"vllm version: {vllm.__version__}", "\n", "=" * 20) # Get paths from args video_path = getattr(args, "video_path", None) @@ -259,6 +318,10 @@ def main(args): num_frames=getattr(args, "num_frames", 16), sampling_rate=getattr(args, "sampling_rate", 16000), ) + elif args.query_type == "multi_audios": + query_result = query_func() + elif args.query_type == "use_audio_in_video": + query_result = query_func() else: query_result = query_func() @@ -270,7 +333,7 @@ def main(args): ) thinker_sampling_params = SamplingParams( - temperature=0.4, + temperature=0.9, top_p=0.9, top_k=-1, max_tokens=1200, @@ -331,6 +394,8 @@ def main(args): total_requests = len(prompts) processed_count = 0 + print(f"query type: {args.query_type}") + for stage_outputs in omni_generator: if stage_outputs.final_output_type == "text": for output in stage_outputs.request_output: @@ -385,7 +450,7 @@ def parse_args(): "--query-type", "-q", type=str, - default="mixed_modalities", + default="use_mixed_modalities", choices=query_map.keys(), help="Query type.", ) 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 0d5020635..e4768a8aa 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 @@ -305,6 +305,26 @@ def get_multi_audios_query(custom_prompt: str | None = None): } +def get_use_audio_in_video_query( + video_path: str | None = None, + audio_path: str | None = None, + custom_prompt: str | None = None, +): + question = custom_prompt or ( + "Describe the content of the video in details, then convert what the baby say into text." + ) + video_url = get_video_url_from_path(video_path) + audio_url = get_audio_url_from_path(audio_path) + return { + "role": "user", + "content": [ + {"type": "video_url", "video_url": {"url": video_url}}, + {"type": "audio_url", "audio_url": {"url": audio_url}}, + {"type": "text", "text": question}, + ], + } + + query_map = { "text": get_text_query, "use_audio": get_audio_query, @@ -312,6 +332,7 @@ query_map = { "use_video": get_video_query, "use_mixed_modalities": get_mixed_modalities_query, "use_multi_audios": get_multi_audios_query, + "use_audio_in_video": get_use_audio_in_video_query, } @@ -372,6 +393,12 @@ def run_multimodal_generation(args) -> None: prompt = query_func(audio_path=audio_path, custom_prompt=custom_prompt) elif args.query_type == "text": prompt = query_func(custom_prompt=custom_prompt) + elif args.query_type == "use_audio_in_video": + prompt = query_func( + video_path=video_path, + audio_path=audio_path, + custom_prompt=custom_prompt, + ) else: prompt = query_func() diff --git a/mkdocs.yml b/mkdocs.yml index 1e8e38f51..15a07450c 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -90,6 +90,7 @@ plugins: api_root_uri: "api" nav_item_prefix: "" # No prefix in navigation tree (clean names) show_full_namespace: false # Show only module name, not full path + on_implicit_namespace_package: skip # Skip directories without __init__.py (e.g., assets) exclude: - "re:vllm_omni\\._.*" # Internal modules - "vllm_omni.diffusion.models.qwen_image" # avoid importing vllm in mkdocs building diff --git a/pyproject.toml b/pyproject.toml index 16c8e3c80..27cdac12d 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,6 +39,7 @@ dependencies = [ "cache-dit==1.2.0", "tqdm>=4.66.0", "torchsde>=0.2.6", # Required for Stable Audio scheduler + "openai-whisper>=20250625", # "vllm==0.12.0", # TODO: fix the entrypoints overwrite problem ] diff --git a/tests/conftest.py b/tests/conftest.py index 5b21f671b..bf40861fc 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,5 +1,11 @@ -import base64 import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" +# Set CPU device for CI environments without GPU +if "VLLM_TARGET_DEVICE" not in os.environ: + os.environ["VLLM_TARGET_DEVICE"] = "cpu" + +import base64 import socket import subprocess import sys @@ -13,11 +19,33 @@ import torch import whisper import yaml from vllm.logger import init_logger -from vllm.utils import get_open_port +from vllm.utils.network_utils import get_open_port logger = init_logger(__name__) +@pytest.fixture(autouse=True) +def default_vllm_config(): + """Set a default VllmConfig for all tests. + + This fixture is auto-used for all tests to ensure that any test + that directly instantiates vLLM CustomOps (e.g., RMSNorm, LayerNorm) + or model components has the required VllmConfig context. + + This fixture is required for vLLM 0.14.0+ where CustomOp initialization + requires a VllmConfig context set via set_current_vllm_config(). + """ + from vllm.config import DeviceConfig, VllmConfig, set_current_vllm_config + + # Use CPU device if no GPU is available (e.g., in CI environments) + has_gpu = torch.cuda.is_available() and torch.cuda.device_count() > 0 + device = "cuda" if has_gpu else "cpu" + device_config = DeviceConfig(device=device) + + with set_current_vllm_config(VllmConfig(device_config=device_config)): + yield + + @pytest.fixture(autouse=True) def clean_gpu_memory_between_tests(): if os.getenv("VLLM_TEST_CLEAN_GPU_MEMORY", "0") != "1": diff --git a/tests/diffusion/attention/test_sequence_parallel.py b/tests/diffusion/attention/test_sequence_parallel.py index 4304de77d..c0061daf9 100644 --- a/tests/diffusion/attention/test_sequence_parallel.py +++ b/tests/diffusion/attention/test_sequence_parallel.py @@ -170,6 +170,11 @@ def test_sequence_parallel( """Test Ulysses attention by comparing with and without SP enabled.""" sequence_parallel_size = ulysses_degree * ring_degree + # Skip if not enough GPUs available + available_gpus = torch_device.device_count() + if available_gpus < sequence_parallel_size: + pytest.skip(f"Test requires {sequence_parallel_size} GPUs but only {available_gpus} available") + # Create temporary files to share results between processes with tempfile.NamedTemporaryFile(delete=False, suffix=".pkl") as f: baseline_output_file = f.name diff --git a/tests/diffusion/distributed/test_comm.py b/tests/diffusion/distributed/test_comm.py index 4a8451696..0e76f27b7 100644 --- a/tests/diffusion/distributed/test_comm.py +++ b/tests/diffusion/distributed/test_comm.py @@ -48,6 +48,11 @@ def test_4d_identity( use_sync: bool, ): """Test that two consecutive all-to-all operations return the original input.""" + # Skip if not enough GPUs available + available_gpus = torch_device.device_count() + if available_gpus < world_size: + pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available") + # Ensure num_heads is divisible by world_size if num_heads % world_size != 0: pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})") @@ -177,6 +182,11 @@ def test_5d_identity( use_sync: bool, ): """Test that two consecutive all-to-all operations return the original input.""" + # Skip if not enough GPUs available + available_gpus = torch_device.device_count() + if available_gpus < world_size: + pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available") + # Ensure num_heads is divisible by world_size if num_heads % world_size != 0: pytest.skip(f"num_heads ({num_heads}) not divisible by world_size ({world_size})") @@ -305,6 +315,11 @@ def test_ring_p2p( head_size: int, ): """Test Ring P2P communication (send_recv).""" + # Skip if not enough GPUs available + available_gpus = torch_device.device_count() + if available_gpus < world_size: + pytest.skip(f"Test requires {world_size} GPUs but only {available_gpus} available") + torch.multiprocessing.spawn( _test_ring_p2p_worker, args=(world_size, dtype, batch_size, num_heads, head_size), diff --git a/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml index 96e9d7fa7..eb9109adc 100644 --- a/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/qwen2_5_omni_ci.yaml @@ -78,6 +78,7 @@ stage_args: trust_remote_code: true enable_prefix_caching: false engine_output_type: audio + max_num_batched_tokens: 4069 engine_input_source: [1] final_output: true final_output_type: audio 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 f8eb59973..0780420fd 100644 --- a/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/qwen3_omni_ci.yaml @@ -45,7 +45,7 @@ stage_args: worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.3 - enforce_eager: false + enforce_eager: true trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav # tensor_parallel_size: 2 @@ -84,6 +84,7 @@ stage_args: max_num_batched_tokens: 1000000 hf_config_name: thinker_config load_format: dummy + async_scheduling: false engine_input_source: [1] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav final_output: true diff --git a/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml index 5e3be865e..0780420fd 100644 --- a/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml +++ b/tests/e2e/offline_inference/stage_configs/rocm/qwen3_omni_ci.yaml @@ -84,6 +84,7 @@ stage_args: max_num_batched_tokens: 1000000 hf_config_name: thinker_config load_format: dummy + async_scheduling: false engine_input_source: [1] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav final_output: true diff --git a/tests/e2e/offline_inference/test_qwen2_5_omni.py b/tests/e2e/offline_inference/test_qwen2_5_omni.py index 9bcc48827..09b0306c1 100644 --- a/tests/e2e/offline_inference/test_qwen2_5_omni.py +++ b/tests/e2e/offline_inference/test_qwen2_5_omni.py @@ -35,7 +35,7 @@ test_params = [(model, stage_config) for model in models] @pytest.mark.core_model @pytest.mark.parametrize("test_config", test_params) -@create_new_process_for_each_test() +@create_new_process_for_each_test("spawn") def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None: """Test processing audio, image, and video together, generating audio output.""" model, stage_config_path = test_config @@ -95,7 +95,7 @@ def test_mixed_modalities_to_audio(omni_runner: type[OmniRunner], test_config: t @pytest.mark.core_model @pytest.mark.parametrize("test_config", test_params) -@create_new_process_for_each_test() +@create_new_process_for_each_test("spawn") def test_mixed_modalities_to_text_only(omni_runner: type[OmniRunner], test_config: tuple[str, str]) -> None: """Test processing audio, image, and video together, generating audio output.""" model, stage_config_path = test_config diff --git a/tests/e2e/offline_inference/test_qwen3_omni.py b/tests/e2e/offline_inference/test_qwen3_omni.py index 71eac2a68..14d117457 100644 --- a/tests/e2e/offline_inference/test_qwen3_omni.py +++ b/tests/e2e/offline_inference/test_qwen3_omni.py @@ -5,6 +5,8 @@ E2E offline tests for Omni model with video input and audio output. """ import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" from pathlib import Path import pytest @@ -14,8 +16,6 @@ from vllm_omni.utils import is_rocm from .conftest import OmniRunner -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] # CI stage config for 2xH100-80G GPUs or AMD GPU MI325 diff --git a/tests/e2e/offline_inference/test_sequence_parallel.py b/tests/e2e/offline_inference/test_sequence_parallel.py index ab30eab5a..29172b7d1 100644 --- a/tests/e2e/offline_inference/test_sequence_parallel.py +++ b/tests/e2e/offline_inference/test_sequence_parallel.py @@ -30,8 +30,6 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.diffusion.distributed.parallel_state import device_count from vllm_omni.diffusion.envs import get_device_name -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" - # Use random weights model for testing models = ["riverclouds/qwen_image_random"] diff --git a/tests/e2e/offline_inference/test_stable_audio_model.py b/tests/e2e/offline_inference/test_stable_audio_model.py index c0201778e..648562573 100644 --- a/tests/e2e/offline_inference/test_stable_audio_model.py +++ b/tests/e2e/offline_inference/test_stable_audio_model.py @@ -1,4 +1,3 @@ -import os import sys from pathlib import Path @@ -15,8 +14,6 @@ if str(REPO_ROOT) not in sys.path: from vllm_omni import Omni -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" - # Use random weights model for CI testing (small, no authentication required) models = ["linyueqian/stable_audio_random"] diff --git a/tests/e2e/offline_inference/test_t2v_model.py b/tests/e2e/offline_inference/test_t2v_model.py index 1dc2e1ff1..0263310e3 100644 --- a/tests/e2e/offline_inference/test_t2v_model.py +++ b/tests/e2e/offline_inference/test_t2v_model.py @@ -11,8 +11,10 @@ if str(REPO_ROOT) not in sys.path: sys.path.insert(0, str(REPO_ROOT)) from vllm_omni import Omni +from vllm_omni.outputs import OmniRequestOutput -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" +# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" models = ["Wan-AI/Wan2.2-T2V-A14B-Diffusers"] @@ -45,10 +47,10 @@ def test_video_diffusion_model(model_name: str): raise ValueError("No request_output found in OmniRequestOutput") req_out = first_output.request_output[0] - if not isinstance(req_out, dict) or "images" not in req_out: + if not isinstance(req_out, OmniRequestOutput) or not hasattr(req_out, "images"): raise ValueError("Invalid request_output structure or missing 'images' key") - frames = req_out["images"][0] + frames = req_out.images[0] assert frames is not None assert hasattr(frames, "shape") diff --git a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py index 606869922..f4fb4371e 100644 --- a/tests/e2e/offline_inference/test_zimage_tensor_parallel.py +++ b/tests/e2e/offline_inference/test_zimage_tensor_parallel.py @@ -23,7 +23,8 @@ from vllm_omni.diffusion.data import DiffusionParallelConfig from vllm_omni.outputs import OmniRequestOutput from vllm_omni.utils.platform_utils import is_npu, is_rocm -os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" +# os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" PROMPT = "a photo of a cat sitting on a laptop keyboard" @@ -129,8 +130,8 @@ def test_zimage_tensor_parallel_tp2(tmp_path: Path): if not torch.cuda.is_available() or torch.cuda.device_count() < 2: pytest.skip("Z-Image TP=2 requires >= 2 CUDA devices.") - height = 256 - width = 256 + height = 512 + width = 512 num_inference_steps = 2 seed = 42 diff --git a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml index 3f727e0d9..c5987ebba 100644 --- a/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/online_serving/stage_configs/qwen3_omni_ci.yaml @@ -77,6 +77,7 @@ stage_args: scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true + async_scheduling: false enable_prefix_caching: false engine_output_type: audio # Final output: audio waveform gpu_memory_utilization: 0.1 diff --git a/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml b/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml index 7420e2745..ef0bae0fa 100644 --- a/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml +++ b/tests/e2e/online_serving/stage_configs/rocm/qwen3_omni_ci.yaml @@ -84,6 +84,7 @@ stage_args: max_num_batched_tokens: 1000000 hf_config_name: thinker_config load_format: dummy + async_scheduling: false engine_input_source: [1] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav final_output: true diff --git a/tests/e2e/online_serving/test_async_omni.py b/tests/e2e/online_serving/test_async_omni.py index a572092b5..112b68ff4 100644 --- a/tests/e2e/online_serving/test_async_omni.py +++ b/tests/e2e/online_serving/test_async_omni.py @@ -1,4 +1,6 @@ import asyncio +import os +import sys from contextlib import ExitStack from pathlib import Path @@ -7,6 +9,8 @@ from vllm.inputs import PromptType from vllm_omni.entrypoints.async_omni import AsyncOmni, ClientRequestState +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + SEED = 42 stage_config = str(Path(__file__).parent / "stage_configs" / "qwen3_omni_thinker_ci.yaml") @@ -55,13 +59,19 @@ async def generate( @pytest.mark.asyncio async def test_abort(): with ExitStack() as after: - engine = AsyncOmni(model=model, stage_configs_path=stage_config) + # Avoid SHM IPC in tests to prevent /dev/shm exhaustion and SIGBUS. + engine = AsyncOmni( + model=model, + stage_configs_path=stage_config, + shm_threshold_bytes=sys.maxsize, + ) after.callback(engine.shutdown) - NUM_REQUESTS = 5 - NUM_EXPECTED_TOKENS = 100 - NUM_EXPECTED_TOKENS_LONG = 1000 - REQUEST_IDS_TO_ABORT = [1, 2, 3] + # Keep token counts modest to reduce flakiness on slow test hardware. + NUM_REQUESTS = 3 + NUM_EXPECTED_TOKENS = 64 + NUM_EXPECTED_TOKENS_LONG = 256 + REQUEST_IDS_TO_ABORT = [1] prompt = "Hello my name is Robert and " @@ -74,8 +84,10 @@ async def test_abort(): tasks.append(asyncio.create_task(generate(engine, request_id, prompt, max_tokens))) # API server cancels requests when they disconnect. + # Explicitly abort in the engine to avoid orphaned requests hanging. for idx in REQUEST_IDS_TO_ABORT: tasks[idx].cancel() + await engine.abort(request_ids[idx]) await asyncio.sleep(0.1) # Confirm the other requests are okay. @@ -83,10 +95,10 @@ async def test_abort(): # Confirm that it was actually canceled. if idx in REQUEST_IDS_TO_ABORT: with pytest.raises((asyncio.CancelledError, GeneratorExit)): - await task + await asyncio.wait_for(task, timeout=60) else: # Otherwise, make sure the request was not impacted. - num_generated_tokens, request_id = await task + num_generated_tokens, request_id = await asyncio.wait_for(task, timeout=180) expected_tokens = NUM_EXPECTED_TOKENS assert num_generated_tokens == expected_tokens, ( f"{request_id} generated {num_generated_tokens} but expected {expected_tokens}" @@ -116,11 +128,16 @@ async def test_build_and_log_summary(monkeypatch): monkeypatch.setattr("vllm_omni.entrypoints.client_request_state.ClientRequestState", MockCRS) with ExitStack() as after: - engine = AsyncOmni(model=model, stage_configs_path=stage_config) + # Avoid SHM IPC in tests to prevent /dev/shm exhaustion and SIGBUS. + engine = AsyncOmni( + model=model, + stage_configs_path=stage_config, + shm_threshold_bytes=sys.maxsize, + ) after.callback(engine.shutdown) prompt = "Hello my name is Robert and " - NUM_EXPECTED_TOKENS = 100 - NUM_REQUESTS = 5 + NUM_EXPECTED_TOKENS = 64 + NUM_REQUESTS = 3 request_ids = [f"request-{i}" for i in range(NUM_REQUESTS)] # Create concurrent requests. diff --git a/tests/e2e/online_serving/test_i2i_multi_image_input.py b/tests/e2e/online_serving/test_i2i_multi_image_input.py index dbaf9f53f..b1fcac1de 100644 --- a/tests/e2e/online_serving/test_i2i_multi_image_input.py +++ b/tests/e2e/online_serving/test_i2i_multi_image_input.py @@ -17,7 +17,7 @@ import openai import pytest from PIL import Image from vllm.assets.image import ImageAsset -from vllm.utils import get_open_port +from vllm.utils.network_utils import get_open_port os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" diff --git a/tests/e2e/online_serving/test_qwen3_omni.py b/tests/e2e/online_serving/test_qwen3_omni.py index b06d58ea1..33a05c1ce 100644 --- a/tests/e2e/online_serving/test_qwen3_omni.py +++ b/tests/e2e/online_serving/test_qwen3_omni.py @@ -4,10 +4,13 @@ E2E Online tests for Qwen3-Omni model with video input and audio output. """ +import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" + import base64 import concurrent.futures import ctypes -import os import signal import socket import subprocess @@ -18,12 +21,10 @@ from pathlib import Path import openai import pytest from vllm.assets.video import VideoAsset -from vllm.utils import get_open_port +from vllm.utils.network_utils import get_open_port from vllm_omni.utils import is_rocm -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] # CI stage config for 2xH100-80G GPUs or AMD GPU MI325 diff --git a/tests/e2e/online_serving/test_qwen3_omni_expansion.py b/tests/e2e/online_serving/test_qwen3_omni_expansion.py index 6a47e96f8..ef10f52fc 100644 --- a/tests/e2e/online_serving/test_qwen3_omni_expansion.py +++ b/tests/e2e/online_serving/test_qwen3_omni_expansion.py @@ -6,6 +6,8 @@ E2E Online tests for Qwen3-Omni model. import concurrent.futures import os + +os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" import time from pathlib import Path @@ -20,8 +22,6 @@ from tests.conftest import ( modify_stage_config, ) -os.environ["VLLM_WORKER_MULTIPROC_METHOD"] = "spawn" - models = ["Qwen/Qwen3-Omni-30B-A3B-Instruct"] # CI stage config for 2*H100-80G GPUs diff --git a/tests/e2e/stage_configs/qwen3_omni_ci.yaml b/tests/e2e/stage_configs/qwen3_omni_ci.yaml index 5106b1854..8870f5de0 100644 --- a/tests/e2e/stage_configs/qwen3_omni_ci.yaml +++ b/tests/e2e/stage_configs/qwen3_omni_ci.yaml @@ -81,6 +81,7 @@ stage_args: distributed_executor_backend: "mp" max_num_batched_tokens: 1000000 hf_config_name: thinker_config + async_scheduling: false engine_input_source: [1] custom_process_input_func: vllm_omni.model_executor.stage_input_processors.qwen3_omni.talker2code2wav final_output: true diff --git a/tests/worker/test_omni_gpu_model_runner.py b/tests/worker/test_omni_gpu_model_runner.py index b0132306c..a600e350a 100644 --- a/tests/worker/test_omni_gpu_model_runner.py +++ b/tests/worker/test_omni_gpu_model_runner.py @@ -69,7 +69,7 @@ def _make_runner(req_ids=("r1", "r2"), hidden_size=4): # Provide a minimal implementation that returns the expected 4-tuple. def _determine_batch_execution_and_padding(**kwargs): - return None, object(), None, None + return None, object(), None, None, None runner._determine_batch_execution_and_padding = _determine_batch_execution_and_padding diff --git a/vllm_omni/assets/video.py b/vllm_omni/assets/video.py new file mode 100644 index 000000000..98b1f7e4e --- /dev/null +++ b/vllm_omni/assets/video.py @@ -0,0 +1,16 @@ +import librosa +import numpy as np +from vllm.assets.video import VideoAsset + + +def extract_video_audio(path: str = None, sampling_rate: int = 16000) -> np.ndarray: + """This function extracts the audio from a video file path and returns the audio as a numpy array. + Args: + path: The path to the video file. + Returns: + The audio as a numpy array. + """ + if not path: + path = VideoAsset(name="baby_reading").video_path + audio_signal, sr = librosa.load(path, sr=sampling_rate) + return audio_signal diff --git a/vllm_omni/config/model.py b/vllm_omni/config/model.py index e074689c9..592c501a6 100644 --- a/vllm_omni/config/model.py +++ b/vllm_omni/config/model.py @@ -1,20 +1,12 @@ import warnings -from importlib.util import find_spec from typing import Any import torch -import vllm.envs as envs 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, _get_and_verify_dtype, get_served_model_name, ) @@ -29,12 +21,11 @@ from vllm.transformers_utils.config import ( get_pooling_config, ) from vllm.transformers_utils.gguf_utils import ( + is_gguf, maybe_patch_hf_config_from_gguf, ) -from vllm.transformers_utils.utils import ( - is_gguf, - maybe_model_redirect, -) +from vllm.transformers_utils.utils import maybe_model_redirect +from vllm.v1.attention.backends.registry import AttentionBackendEnum import vllm_omni.model_executor.models as me_models @@ -146,14 +137,6 @@ class OmniModelConfig(ModelConfig): 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 " # 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", @@ -183,111 +166,16 @@ class OmniModelConfig(ModelConfig): 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() - # Try to load image processor config, but allow it to fail for stages that don't need it - try: - self.hf_image_processor_config = get_hf_image_processor_config( - self.model, hf_token=self.hf_token, revision=self.revision - ) - except (OSError, ValueError, IndexError) as e: - # Some stages (e.g., code2wav, talker) don't need image processor - # Log warning but allow initialization to continue - logger.warning( - f"Failed to load image processor config for model '{self.model}': {e}. " - "This is expected for stages that don't require image processing." - ) - self.hf_image_processor_config = None + self.hf_image_processor_config = get_hf_image_processor_config( + self.model, hf_token=self.hf_token, revision=self.revision + ) + self.model_arch_config = self.get_model_arch_config() architectures = self.architectures registry = self.registry is_generative_model = registry.is_text_generation_model(architectures, self) is_pooling_model = registry.is_pooling_model(architectures, self) - def _task_to_convert(task: TaskOption) -> ConvertType: - if task == "embedding" or task == "embed": - return "embed" - if task == "classify": - return "classify" - if task == "reward": - return "reward" - if task == "score": - new_task = self._get_default_pooling_task(architectures) - return "classify" if new_task == "classify" else "embed" - - return "none" - - if self.task is not None: - runner: RunnerOption = "auto" - convert: ConvertOption = "auto" - msg_prefix = ( - "The 'task' option has been deprecated and will be removed in v0.13.0 or v1.0, whichever comes first." - ) - msg_hint = "Please remove this option." - - is_generative_task = self.task in _RUNNER_TASKS["generate"] - is_pooling_task = self.task in _RUNNER_TASKS["pooling"] - - if is_generative_model and is_pooling_model: - if is_generative_task: - runner = "generate" - convert = "auto" - msg_hint = ( - "Please replace this option with `--runner " - "generate` to continue using this model " - "as a generative model." - ) - elif is_pooling_task: - runner = "pooling" - convert = "auto" - msg_hint = ( - "Please replace this option with `--runner " - "pooling` to continue using this model " - "as a pooling model." - ) - else: # task == "auto" - pass - elif is_generative_model or is_pooling_model: - if is_generative_task: - runner = "generate" - convert = "auto" - msg_hint = "Please remove this option" - elif is_pooling_task: - runner = "pooling" - convert = _task_to_convert(self.task) - msg_hint = ( - "Please replace this option with `--convert " - f"{convert}` to continue using this model " - "as a pooling model." - ) - else: # task == "auto" - pass - else: - # 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 - - msg = f"{msg_prefix} {msg_hint}" - warnings.warn(msg, DeprecationWarning, stacklevel=2) - self.runner_type = self._get_runner_type(architectures, self.runner) self.convert_type = self._get_convert_type(architectures, self.runner_type, self.convert) @@ -325,9 +213,12 @@ class OmniModelConfig(ModelConfig): if getattr(self.pooler_config, k) is None: setattr(self.pooler_config, k, v) - default_pooling_type = self._model_info.default_pooling_type - if self.pooler_config.pooling_type is None: - self.pooler_config.pooling_type = default_pooling_type + default_seq_pooling_type = self._model_info.default_seq_pooling_type + if self.pooler_config.seq_pooling_type is None: + self.pooler_config.seq_pooling_type = default_seq_pooling_type + default_tok_pooling_type = self._model_info.default_tok_pooling_type + if self.pooler_config.tok_pooling_type is None: + self.pooler_config.tok_pooling_type = default_tok_pooling_type self.dtype: torch.dtype = _get_and_verify_dtype( self.model, @@ -339,6 +230,11 @@ class OmniModelConfig(ModelConfig): self.original_max_model_len = self.max_model_len self.max_model_len = self.get_and_verify_max_len(self.max_model_len) + + if self.is_encoder_decoder: + self.mm_processor_cache_gb = 0 + logger.info("Encoder-decoder model detected, disabling mm processor cache.") + # Init multimodal config if needed if self._model_info.supports_multimodal: if mm_encoder_tp_mode == "data" and not self._model_info.supports_multimodal_encoder_tp_data: @@ -382,7 +278,7 @@ class OmniModelConfig(ModelConfig): # Avoid running try_verify_and_update_config multiple times self.config_updated = False - + self._try_verify_and_update_model_config() self._verify_quantization() self._verify_cuda_graph() self._verify_bnb_config() diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 7918e1687..3c262b7c9 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -3,6 +3,8 @@ from __future__ import annotations from collections import defaultdict from time import time +import numpy as np +from vllm.compilation.cuda_graph import CUDAGraphStat from vllm.distributed.kv_events import KVEventBatch from vllm.distributed.kv_transfer.kv_connector.v1.metrics import KVConnectorStats from vllm.logger import init_logger @@ -10,10 +12,13 @@ from vllm.v1.core.sched.output import SchedulerOutput from vllm.v1.core.sched.scheduler import Scheduler as VLLMScheduler from vllm.v1.core.sched.utils import check_stop, remove_all from vllm.v1.engine import EngineCoreOutput, EngineCoreOutputs +from vllm.v1.metrics.perf import PerfStats from vllm.v1.outputs import ModelRunnerOutput from vllm.v1.request import Request, RequestStatus from vllm.v1.spec_decode.metrics import SpecDecodingStats +logger = init_logger(__name__) + class OmniARScheduler(VLLMScheduler): """ @@ -73,6 +78,11 @@ class OmniARScheduler(VLLMScheduler): pooler_outputs = model_runner_output.pooler_output num_nans_in_logits = model_runner_output.num_nans_in_logits kv_connector_output = model_runner_output.kv_connector_output + cudagraph_stats: CUDAGraphStat | None = model_runner_output.cudagraph_stats + + perf_stats: PerfStats | None = None + if self.perf_metrics and self.perf_metrics.is_enabled(): + perf_stats = self.perf_metrics.get_step_perf_stats_per_gpu(scheduler_output) outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None @@ -131,11 +141,14 @@ class OmniARScheduler(VLLMScheduler): spec_decoding_stats, num_draft_tokens=num_draft_tokens, num_accepted_tokens=num_accepted, + num_invalid_spec_tokens=scheduler_output.num_invalid_spec_tokens, + request_id=req_id, ) stopped = False new_logprobs = None new_token_ids = generated_token_ids + pooler_output = pooler_outputs[req_index] if pooler_outputs else None kv_transfer_params = None status_before_stop = request.status @@ -143,14 +156,31 @@ 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] + if pooler_output: + # Note: For multimodal outputs, we skip intermediate stop checks. if request.output_token_ids: - stopped = check_stop(request, self.max_model_len, pooler_output) - + stopped = check_stop(request, self.max_model_len) + routed_experts = None if stopped: + if self.vllm_config.model_config.enable_return_routed_experts: + kv_blocks = self.kv_cache_manager.get_blocks(request.request_id) + block_ids = kv_blocks.get_block_ids()[0] + num_tokens = request.num_tokens - 1 + + # compute slot mapping + block_ids_array = np.array(block_ids, dtype=np.int32) + num_blocks = len(block_ids) + block_size = self.block_size + + # generate block offsets + block_offsets = np.arange(0, block_size) + + # compute slot mapping: slot = block_id * block_size + offset + slot_mapping = ( + block_offsets.reshape((1, block_size)) + block_ids_array.reshape((num_blocks, 1)) * block_size + ).flatten()[:num_tokens] + + routed_experts = self.routed_experts_reader.get_routed_experts(indices=slot_mapping) kv_transfer_params = self._free_request(request) if status_before_stop == RequestStatus.RUNNING: stopped_running_reqs.add(request) @@ -165,7 +195,13 @@ class OmniARScheduler(VLLMScheduler): 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) + ok = struct_output_request.grammar.accept_tokens(req_id, new_token_ids) + if not ok: + logger.warning( + "Unexpected: grammar rejected tokens %s for request %s.", + new_token_ids, + req_id, + ) 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] @@ -187,6 +223,7 @@ class OmniARScheduler(VLLMScheduler): kv_transfer_params=kv_transfer_params, trace_headers=request.trace_headers, num_cached_tokens=request.num_cached_tokens, + routed_experts=routed_experts, num_nans_in_logits=request.num_nans_in_logits, ) ) @@ -201,6 +238,20 @@ class OmniARScheduler(VLLMScheduler): # This is a rare case and unlikely to impact performance. self.waiting.remove_requests(stopped_preempted_reqs) + if failed_kv_load_req_ids and not self.recompute_kv_load_failures: + requests = [self.requests[req_id] for req_id in failed_kv_load_req_ids] + self.finish_requests(failed_kv_load_req_ids, RequestStatus.FINISHED_ERROR) + for request in requests: + outputs[request.client_index].append( + EngineCoreOutput( + request_id=request.request_id, + new_token_ids=[], + finish_reason=request.get_finished_reason(), + events=request.take_events(), + trace_headers=request.trace_headers, + num_cached_tokens=request.num_cached_tokens, + ) + ) # KV Connector: update state for finished KV Transfers. if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) @@ -238,7 +289,7 @@ class OmniARScheduler(VLLMScheduler): engine_core_outputs[client_index] = EngineCoreOutputs(finished_requests=finished_set) finished_req_ids.clear() - if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats)) is not None: + if (stats := self.make_stats(spec_decoding_stats, kv_connector_stats, cudagraph_stats, perf_stats)) is not None: # Return stats to only one of the front-ends. if (eco := next(iter(engine_core_outputs.values()), None)) is None: # We must return the stats even if there are no request diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 7acef3bdc..2d04faeee 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -90,7 +90,7 @@ class OmniGenerationScheduler(VLLMScheduler): any_request = self.running[0] num_common_prefix_blocks = self.kv_cache_manager.get_num_common_prefix_blocks(any_request.request_id) - # Assemble SchedulerOutput (align with v0.12.0) + # Assemble SchedulerOutput (align with v0.14.0) if self.use_v2_model_runner: # No resumed reqs in fast path; pass prefill_token_ids for new reqs. new_reqs_data = [ @@ -129,7 +129,7 @@ class OmniGenerationScheduler(VLLMScheduler): preempted_req_ids=set(), ) - # Record the request ids scheduled in this step (v0.12.0 behavior). + # Record the request ids scheduled in this step (v0.14.0 behavior). self.prev_step_scheduled_req_ids.clear() self.prev_step_scheduled_req_ids.update(num_scheduled_tokens.keys()) @@ -176,7 +176,7 @@ 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) + # Merge connector-side stats (align with v0.14.0) if kv_connector_stats and self.connector: kv_stats = self.connector.get_kv_connector_stats() if kv_stats: @@ -294,7 +294,7 @@ class OmniGenerationScheduler(VLLMScheduler): if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) - # Collect and publish KV cache events (align with v0.12.0) + # Collect and publish KV cache events (align with v0.14.0) events = self.kv_cache_manager.take_events() if self.connector is not None: connector_events = self.connector.take_events() diff --git a/vllm_omni/diffusion/forward_context.py b/vllm_omni/diffusion/forward_context.py index 3c33f8105..d898bceaf 100644 --- a/vllm_omni/diffusion/forward_context.py +++ b/vllm_omni/diffusion/forward_context.py @@ -86,5 +86,14 @@ def set_forward_context( attn_metadata=attn_metadata, split_text_embed_in_sp=split_text_embed_in_sp, ) + # vLLM CustomOp dispatch (e.g. QKVParallelLinear) requires a global + # vLLM config set via set_current_vllm_config(). with override_forward_context(forward_context): - yield + if vllm_config is None: + yield + else: + # Local import to avoid importing vllm.config.vllm at module import time. + from vllm.config.vllm import set_current_vllm_config + + with set_current_vllm_config(vllm_config): + yield diff --git a/vllm_omni/diffusion/models/glm_image/__init__.py b/vllm_omni/diffusion/models/glm_image/__init__.py new file mode 100644 index 000000000..fc8256d8d --- /dev/null +++ b/vllm_omni/diffusion/models/glm_image/__init__.py @@ -0,0 +1,21 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""GLM Image diffusion model components.""" + +from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageKVCache, + GlmImageTransformer2DModel, +) +from vllm_omni.diffusion.models.glm_image.pipeline_glm_image import ( + GlmImagePipeline, + get_glm_image_post_process_func, + get_glm_image_pre_process_func, +) + +__all__ = [ + "GlmImageKVCache", + "GlmImagePipeline", + "GlmImageTransformer2DModel", + "get_glm_image_post_process_func", + "get_glm_image_pre_process_func", +] diff --git a/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py new file mode 100644 index 000000000..615b9194a --- /dev/null +++ b/vllm_omni/diffusion/models/glm_image/glm_image_transformer.py @@ -0,0 +1,792 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +from collections.abc import Iterable +from enum import Enum +from typing import Any + +import torch +import torch.nn as nn +from diffusers.models.attention import FeedForward +from diffusers.models.modeling_outputs import Transformer2DModelOutput +from diffusers.models.transformers.transformer_glm_image import GlmImageCombinedTimestepSizeEmbeddings +from vllm.logger import init_logger +from vllm.model_executor.layers.linear import QKVParallelLinear +from vllm.model_executor.model_loader.weight_utils import default_weight_loader + +from vllm_omni.diffusion.attention.layer import Attention +from vllm_omni.diffusion.cache.base import CachedTransformer +from vllm_omni.diffusion.data import OmniDiffusionConfig +from vllm_omni.diffusion.layers.rope import RotaryEmbedding + +logger = init_logger(__name__) + + +class GlmImageImageProjector(nn.Module): + """Projects latent image patches to transformer hidden dimension.""" + + def __init__( + self, + in_channels: int = 16, + hidden_size: int = 2560, + patch_size: int = 2, + ): + super().__init__() + self.patch_size = patch_size + self.proj = nn.Linear(in_channels * patch_size**2, hidden_size) + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + batch_size, channel, height, width = hidden_states.shape + post_patch_height = height // self.patch_size + post_patch_width = width // self.patch_size + + # Reshape: [B, C, H, W] -> [B, H', W', C*p*p] -> [B, H'*W', C*p*p] + hidden_states = hidden_states.reshape( + batch_size, channel, post_patch_height, self.patch_size, post_patch_width, self.patch_size + ) + hidden_states = hidden_states.permute(0, 2, 4, 1, 3, 5).flatten(3, 5).flatten(1, 2) + hidden_states = self.proj(hidden_states) + return hidden_states + + +class GlmImageRotaryPosEmbed(nn.Module): + """Rotary positional embedding for 2D image patches.""" + + def __init__(self, dim: int, patch_size: int, theta: float = 10000.0) -> None: + super().__init__() + self.dim = dim + self.patch_size = patch_size + self.theta = theta + + def forward(self, hidden_states: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]: + batch_size, num_channels, height, width = hidden_states.shape + height, width = height // self.patch_size, width // self.patch_size + + dim_h, dim_w = self.dim // 2, self.dim // 2 + h_inv_freq = 1.0 / ( + self.theta ** (torch.arange(0, dim_h, 2, dtype=torch.float32)[: (dim_h // 2)].float() / dim_h) + ) + w_inv_freq = 1.0 / ( + self.theta ** (torch.arange(0, dim_w, 2, dtype=torch.float32)[: (dim_w // 2)].float() / dim_w) + ) + h_seq = torch.arange(height, device=hidden_states.device) + w_seq = torch.arange(width, device=hidden_states.device) + h_inv_freq = h_inv_freq.to(hidden_states.device) + w_inv_freq = w_inv_freq.to(hidden_states.device) + + freqs_h = torch.outer(h_seq, h_inv_freq) + freqs_w = torch.outer(w_seq, w_inv_freq) + + # Create position matrices: [height, 1, dim//4] and [1, width, dim//4] + freqs_h = freqs_h.unsqueeze(1).expand(height, width, -1) + freqs_w = freqs_w.unsqueeze(0).expand(height, width, -1) + + # Concatenate: [height, width, dim//2] -> [height, width, dim] + freqs = torch.cat([freqs_h, freqs_w], dim=-1) + freqs = torch.cat([freqs, freqs], dim=-1) + freqs = freqs.reshape(height * width, -1) + return (freqs.cos(), freqs.sin()) + + +class GlmImageAdaLayerNormZero(nn.Module): + """Adaptive LayerNorm with zero initialization for both image and text streams.""" + + def __init__(self, embedding_dim: int, dim: int) -> None: + super().__init__() + self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.norm_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.linear = nn.Linear(embedding_dim, 12 * dim, bias=True) + + def forward( + self, hidden_states: torch.Tensor, encoder_hidden_states: torch.Tensor, temb: torch.Tensor + ) -> tuple[torch.Tensor, ...]: + dtype = hidden_states.dtype + norm_hidden_states = self.norm(hidden_states).to(dtype=dtype) + norm_encoder_hidden_states = self.norm_context(encoder_hidden_states).to(dtype=dtype) + + emb = self.linear(temb) + ( + shift_msa, + c_shift_msa, + scale_msa, + c_scale_msa, + gate_msa, + c_gate_msa, + shift_mlp, + c_shift_mlp, + scale_mlp, + c_scale_mlp, + gate_mlp, + c_gate_mlp, + ) = emb.chunk(12, dim=1) + + hidden_states = norm_hidden_states * (1 + scale_msa.unsqueeze(1)) + shift_msa.unsqueeze(1) + encoder_hidden_states = norm_encoder_hidden_states * (1 + c_scale_msa.unsqueeze(1)) + c_shift_msa.unsqueeze(1) + + return ( + hidden_states, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + encoder_hidden_states, + c_gate_msa, + c_shift_mlp, + c_scale_mlp, + c_gate_mlp, + ) + + +class GlmImageAdaLayerNormContinuous(nn.Module): + """Final AdaLN for output projection (no activation before Linear).""" + + def __init__( + self, + embedding_dim: int, + conditioning_embedding_dim: int, + elementwise_affine: bool = True, + eps: float = 1e-5, + bias: bool = True, + ): + super().__init__() + self.linear = nn.Linear(conditioning_embedding_dim, embedding_dim * 2, bias=bias) + self.norm = nn.LayerNorm(embedding_dim, eps=eps, elementwise_affine=elementwise_affine) + + def forward(self, x: torch.Tensor, conditioning_embedding: torch.Tensor) -> torch.Tensor: + # NO SiLU here + emb = self.linear(conditioning_embedding.to(x.dtype)) + scale, shift = torch.chunk(emb, 2, dim=1) + x = self.norm(x) * (1 + scale)[:, None, :] + shift[:, None, :] + return x + + +class KVCacheMode(Enum): + """Mode for KV cache operations. + + - WRITE: Store the K/V tensors from condition images + - READ: Concatenate cached K/V with current K/V + - SKIP: Do not use cache (pass-through) + """ + + WRITE = "write" + READ = "read" + SKIP = "skip" + + +class GlmImageLayerKVCache: + """KV cache for a single attention layer. + + Stores key and value tensors for image editing. The cache accumulates + KV pairs during write mode and provides them during read mode. + + Shape convention (vllm-omni): + key/value: [batch_size, seq_length, num_heads, head_dim] + """ + + def __init__(self): + self.k_cache: torch.Tensor | None = None + self.v_cache: torch.Tensor | None = None + + def store(self, key: torch.Tensor, value: torch.Tensor) -> None: + """Store or accumulate KV tensors. + + If cache is empty, stores the tensors directly. + If cache is not empty, concatenates new tensors along seq_length dim. + + Args: + key: Key tensor of shape [B, S, H, D] + value: Value tensor of shape [B, S, H, D] + """ + if self.k_cache is None: + self.k_cache = key + self.v_cache = value + else: + # Concatenate along sequence dimension (dim=1 for [B, S, H, D]) + self.k_cache = torch.cat([self.k_cache, key], dim=1) + self.v_cache = torch.cat([self.v_cache, value], dim=1) + + def get(self) -> tuple[torch.Tensor | None, torch.Tensor | None]: + """Get cached KV tensors. + + Returns: + Tuple of (k_cache, v_cache), both may be None if cache is empty. + """ + return self.k_cache, self.v_cache + + def clear(self) -> None: + """Clear the cache.""" + self.k_cache = None + self.v_cache = None + + @property + def is_empty(self) -> bool: + """Check if cache is empty.""" + return self.k_cache is None + + def __repr__(self) -> str: + if self.is_empty: + return "GlmImageLayerKVCache(empty)" + return f"GlmImageLayerKVCache(k_shape={self.k_cache.shape}, v_shape={self.v_cache.shape})" + + +class GlmImageKVCache: + """Container for all layers' KV caches. + + Manages KV cache for all transformer layers in GLM-Image model. + Provides a unified interface for setting mode and clearing cache. + + Args: + num_layers: Number of transformer layers in the model. + + Example: + kv_cache = GlmImageKVCache(num_layers=28) + kv_cache.set_mode(KVCacheMode.WRITE) + # ... process condition image ... + kv_cache.set_mode(KVCacheMode.READ) + # ... process target image ... + kv_cache.clear() + """ + + def __init__(self, num_layers: int): + self.num_layers = num_layers + self.caches = [GlmImageLayerKVCache() for _ in range(num_layers)] + self._mode: KVCacheMode | None = None + + def __getitem__(self, layer_idx: int) -> GlmImageLayerKVCache: + """Get cache for a specific layer. + + Args: + layer_idx: Index of the layer (0-indexed). + + Returns: + GlmImageLayerKVCache for the specified layer. + + Raises: + IndexError: If layer_idx is out of range. + """ + if layer_idx < 0 or layer_idx >= self.num_layers: + raise IndexError(f"Layer index {layer_idx} out of range [0, {self.num_layers})") + return self.caches[layer_idx] + + def __len__(self) -> int: + """Return number of layers.""" + return self.num_layers + + @property + def mode(self) -> KVCacheMode | None: + """Get current cache mode.""" + return self._mode + + def set_mode(self, mode: KVCacheMode | str | None) -> None: + """Set cache mode for all layers. + + Args: + mode: Cache mode (WRITE, READ, SKIP) or string ("write", "read", "skip"). + Use None to disable cache operations. + + Raises: + ValueError: If mode is an invalid string. + """ + if mode is None: + self._mode = None + elif isinstance(mode, str): + try: + self._mode = KVCacheMode(mode.lower()) + except ValueError: + raise ValueError(f"Invalid mode: '{mode}', must be one of 'write', 'read', 'skip'") + else: + self._mode = mode + + def clear(self) -> None: + """Clear cache for all layers and reset mode.""" + for cache in self.caches: + cache.clear() + self._mode = None + + @property + def is_empty(self) -> bool: + """Check if all layer caches are empty.""" + return all(cache.is_empty for cache in self.caches) + + def __repr__(self) -> str: + mode_str = self._mode.value if self._mode else "None" + return f"GlmImageKVCache(num_layers={self.num_layers}, mode={mode_str}, is_empty={self.is_empty})" + + +class GlmImageAttention(nn.Module): + """ + Joint attention for GLM-Image model using vllm-omni's optimized attention. + + This combines text and image streams for joint attention computation. + Supports KV caching for image editing workflows via external cache. + """ + + def __init__( + self, + dim: int, + num_heads: int, + head_dim: int, + out_bias: bool = True, + eps: float = 1e-5, + ): + super().__init__() + self.dim = dim + self.num_heads = num_heads + self.head_dim = head_dim + self.inner_dim = num_heads * head_dim + + # QKV projection (fused for efficiency) + self.to_qkv = QKVParallelLinear( + hidden_size=dim, + head_size=head_dim, + total_num_heads=num_heads, + disable_tp=True, + bias=True, + ) + + # QK normalization (LayerNorm, not RMSNorm for GLM-Image) + self.norm_q = nn.LayerNorm(head_dim, elementwise_affine=False, eps=eps) + self.norm_k = nn.LayerNorm(head_dim, elementwise_affine=False, eps=eps) + + # Output projection + self.to_out = nn.Sequential( + nn.Linear(self.inner_dim, dim, bias=out_bias), + nn.Dropout(0.0), + ) + + # RoPE and attention + self.rope = RotaryEmbedding(is_neox_style=False) + self.attn = Attention( + num_heads=num_heads, + head_size=head_dim, + softmax_scale=1.0 / (head_dim**0.5), + causal=False, + ) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + kv_cache: GlmImageLayerKVCache | None = None, + kv_cache_mode: KVCacheMode | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass for joint attention. + + Args: + hidden_states: Image hidden states [B, img_seq_len, D] + encoder_hidden_states: Text hidden states [B, text_seq_len, D] + image_rotary_emb: Tuple of (cos, sin) for RoPE + attention_mask: Optional attention mask for text tokens + kv_cache: Optional layer KV cache for image editing + kv_cache_mode: Cache mode (WRITE, READ, SKIP) + + Returns: + Tuple of (image_hidden_states, text_hidden_states) + """ + dtype = encoder_hidden_states.dtype + batch_size, text_seq_length, _ = encoder_hidden_states.shape + + # Concatenate text and image: [text, image] + hidden_states_combined = torch.cat([encoder_hidden_states, hidden_states], dim=1) + + # QKV projection + qkv, _ = self.to_qkv(hidden_states_combined) + query, key, value = qkv.chunk(3, dim=-1) + + # Reshape: [B, S, H*D] -> [B, S, H, D] + query = query.unflatten(-1, (self.num_heads, -1)) + key = key.unflatten(-1, (self.num_heads, -1)) + value = value.unflatten(-1, (self.num_heads, -1)) + + # QK normalization + query = self.norm_q(query).to(dtype=dtype) + key = self.norm_k(key).to(dtype=dtype) + + # Apply RoPE only to image tokens (not text tokens) + if image_rotary_emb is not None: + cos, sin = image_rotary_emb + cos = cos.to(query.dtype) + sin = sin.to(query.dtype) + # Only apply RoPE to image part (after text_seq_length) + query_img = query[:, text_seq_length:, :, :] + key_img = key[:, text_seq_length:, :, :] + from diffusers.models.embeddings import apply_rotary_emb + + query_img = apply_rotary_emb(query_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) + # key_img = self.rope(key_img, cos, sin) + key_img = apply_rotary_emb(key_img, image_rotary_emb, sequence_dim=1, use_real_unbind_dim=-2) + query = torch.cat([query[:, :text_seq_length, :, :], query_img], dim=1) + key = torch.cat([key[:, :text_seq_length, :, :], key_img], dim=1) + + # Handle KV cache for image editing + if kv_cache is not None and kv_cache_mode is not None: + if kv_cache_mode == KVCacheMode.WRITE: + kv_cache.store(key, value) + elif kv_cache_mode == KVCacheMode.READ: + k_cached, v_cached = kv_cache.get() + if k_cached is not None: + key = torch.cat([k_cached, key], dim=1) + value = torch.cat([v_cached, value], dim=1) + # KVCacheMode.SKIP: do nothing + + # Attention computation + hidden_states_out = self.attn(query, key, value) + hidden_states_out = hidden_states_out.flatten(2, 3) + hidden_states_out = hidden_states_out.to(dtype) + + # Output projection + hidden_states_out = self.to_out(hidden_states_out) + + # Split back to text and image + encoder_hidden_states_out = hidden_states_out[:, :text_seq_length, :] + hidden_states_out = hidden_states_out[:, text_seq_length:, :] + + return hidden_states_out, encoder_hidden_states_out + + +class GlmImageTransformerBlock(nn.Module): + """Single transformer block for GLM-Image.""" + + def __init__( + self, + dim: int = 2560, + num_attention_heads: int = 64, + attention_head_dim: int = 40, + time_embed_dim: int = 512, + ) -> None: + super().__init__() + + # 1. Attention with AdaLN + self.norm1 = GlmImageAdaLayerNormZero(time_embed_dim, dim) + self.attn1 = GlmImageAttention( + dim=dim, + num_heads=num_attention_heads, + head_dim=attention_head_dim, + ) + + # 2. Feedforward + self.norm2 = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.norm2_context = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-5) + self.ff = FeedForward(dim=dim, dim_out=dim, activation_fn="gelu-approximate") + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + temb: torch.Tensor, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + attention_mask: torch.Tensor | None = None, + attention_kwargs: dict[str, Any] | None = None, + kv_cache: GlmImageLayerKVCache | None = None, + kv_cache_mode: KVCacheMode | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + """ + Forward pass for transformer block. + + Args: + hidden_states: Image hidden states + encoder_hidden_states: Text hidden states + temb: Timestep embedding + image_rotary_emb: RoPE embeddings + attention_mask: Text attention mask + attention_kwargs: Additional attention arguments + kv_cache: Layer-specific KV cache for image editing + kv_cache_mode: Cache mode (WRITE, READ, SKIP) + + Returns: + Tuple of (image_hidden_states, text_hidden_states) + """ + # 1. Timestep conditioning via AdaLN + ( + norm_hidden_states, + gate_msa, + shift_mlp, + scale_mlp, + gate_mlp, + norm_encoder_hidden_states, + c_gate_msa, + c_shift_mlp, + c_scale_mlp, + c_gate_mlp, + ) = self.norm1(hidden_states, encoder_hidden_states, temb) + + # 2. Attention + attn_hidden_states, attn_encoder_hidden_states = self.attn1( + hidden_states=norm_hidden_states, + encoder_hidden_states=norm_encoder_hidden_states, + image_rotary_emb=image_rotary_emb, + attention_mask=attention_mask, + kv_cache=kv_cache, + kv_cache_mode=kv_cache_mode, + ) + hidden_states = hidden_states + attn_hidden_states * gate_msa.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + attn_encoder_hidden_states * c_gate_msa.unsqueeze(1) + + # 3. Feedforward + norm_hidden_states = self.norm2(hidden_states) * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1) + norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states) * ( + 1 + c_scale_mlp.unsqueeze(1) + ) + c_shift_mlp.unsqueeze(1) + + ff_output = self.ff(norm_hidden_states) + ff_output_context = self.ff(norm_encoder_hidden_states) + hidden_states = hidden_states + ff_output * gate_mlp.unsqueeze(1) + encoder_hidden_states = encoder_hidden_states + ff_output_context * c_gate_mlp.unsqueeze(1) + + return hidden_states, encoder_hidden_states + + +class GlmImageTransformer2DModel(CachedTransformer): + """ + GLM-Image Transformer model for 2D image generation. + + This is the vllm-omni optimized version of the GLM-Image DiT model. + + Args: + od_config: OmniDiffusionConfig containing model configuration. The + transformer hyper-parameters (e.g. patch size / channels / heads) + are read from `od_config.tf_model_config`. + """ + + def __init__( + self, + od_config: OmniDiffusionConfig, + ): + super().__init__() + + patch_size = od_config.tf_model_config.patch_size + in_channels = od_config.tf_model_config.in_channels + out_channels = od_config.tf_model_config.out_channels + num_attention_heads = od_config.tf_model_config.num_attention_heads + attention_head_dim = od_config.tf_model_config.attention_head_dim + time_embed_dim = od_config.tf_model_config.time_embed_dim + condition_dim = od_config.tf_model_config.condition_dim + prior_vq_quantizer_codebook_size = od_config.tf_model_config.prior_vq_quantizer_codebook_size + text_embed_dim = od_config.tf_model_config.text_embed_dim + + # Get num_layers from config if available + model_config = od_config.tf_model_config + if model_config is not None and hasattr(model_config, "num_layers"): + num_layers = model_config.num_layers + + self.od_config = od_config + self.patch_size = patch_size + self.in_channels = in_channels + self.out_channels = out_channels + + # GlmImage uses 2 additional SDXL-like conditions - target_size, crop_coords + pooled_projection_dim = 2 * 2 * condition_dim + inner_dim = num_attention_heads * attention_head_dim + + # 1. RoPE + self.rope = GlmImageRotaryPosEmbed(attention_head_dim, patch_size, theta=10000.0) + + # 2. Patch & Text-timestep embedding + self.image_projector = GlmImageImageProjector(in_channels, inner_dim, patch_size) + self.glyph_projector = FeedForward(text_embed_dim, inner_dim, inner_dim=inner_dim, activation_fn="gelu") + self.prior_token_embedding = nn.Embedding(prior_vq_quantizer_codebook_size, inner_dim) + self.prior_projector = FeedForward(inner_dim, inner_dim, inner_dim=inner_dim, activation_fn="linear-silu") + + self.time_condition_embed = GlmImageCombinedTimestepSizeEmbeddings( + embedding_dim=time_embed_dim, + condition_dim=condition_dim, + pooled_projection_dim=pooled_projection_dim, + timesteps_dim=time_embed_dim, + ) + + # 3. Transformer blocks + self.transformer_blocks = nn.ModuleList( + [ + GlmImageTransformerBlock(inner_dim, num_attention_heads, attention_head_dim, time_embed_dim) + for _ in range(num_layers) + ] + ) + + # 4. Output projection + self.norm_out = GlmImageAdaLayerNormContinuous(inner_dim, time_embed_dim, elementwise_affine=False) + self.proj_out = nn.Linear(inner_dim, patch_size * patch_size * out_channels, bias=True) + + def forward( + self, + hidden_states: torch.Tensor, + encoder_hidden_states: torch.Tensor, + prior_token_id: torch.Tensor, + prior_token_drop: torch.Tensor, + timestep: torch.LongTensor, + target_size: torch.Tensor, + crop_coords: torch.Tensor, + attention_kwargs: dict[str, Any] | None = None, + return_dict: bool = True, + attention_mask: torch.Tensor | None = None, + image_rotary_emb: tuple[torch.Tensor, torch.Tensor] | None = None, + kv_cache: GlmImageKVCache | None = None, + ) -> torch.Tensor | Transformer2DModelOutput: + """ + Forward pass of the GLM-Image Transformer. + + Args: + hidden_states: Input latent tensor of shape [B, C, H, W]. + encoder_hidden_states: Text embeddings of shape [B, S, D]. + prior_token_id: Prior VQ token IDs. + prior_token_drop: Mask for dropping prior tokens (CFG). + timestep: Diffusion timestep. + target_size: Target image size for conditioning. + crop_coords: Crop coordinates for conditioning. + attention_kwargs: Additional attention arguments. + return_dict: Whether to return a dataclass. + attention_mask: Optional attention mask for text tokens. + image_rotary_emb: Pre-computed rotary embeddings. + kv_cache: Optional KV cache for image editing. When provided, + the cache's mode determines behavior: + - WRITE: Store KV from condition images + - READ: Use cached KV during generation + - SKIP: No caching (same as None) + + Returns: + Output tensor or Transformer2DModelOutput. + """ + batch_size, num_channels, height, width = hidden_states.shape + + # Get KV cache mode + kv_cache_mode = kv_cache.mode if kv_cache is not None else None + + # 1. RoPE + if image_rotary_emb is None: + image_rotary_emb = self.rope(hidden_states) + # Move to correct device + image_rotary_emb = ( + image_rotary_emb[0].to(hidden_states.device), + image_rotary_emb[1].to(hidden_states.device), + ) + + # 2. Patch & Timestep embeddings + p = self.patch_size + post_patch_height = height // p + post_patch_width = width // p + + hidden_states = self.image_projector(hidden_states) + encoder_hidden_states = self.glyph_projector(encoder_hidden_states) + + # Prior embedding with dropout + prior_embedding = self.prior_token_embedding(prior_token_id) + prior_embedding[prior_token_drop] *= 0.0 + prior_hidden_states = self.prior_projector(prior_embedding) + hidden_states = hidden_states + prior_hidden_states + + # Timestep conditioning + temb = self.time_condition_embed(timestep, target_size, crop_coords, hidden_states.dtype) + + # 3. Transformer blocks + for layer_idx, block in enumerate(self.transformer_blocks): + # Get layer-specific KV cache if available + layer_kv_cache = kv_cache[layer_idx] if kv_cache is not None else None + + hidden_states, encoder_hidden_states = block( + hidden_states, + encoder_hidden_states, + temb, + image_rotary_emb, + attention_mask, + attention_kwargs, + kv_cache=layer_kv_cache, + kv_cache_mode=kv_cache_mode, + ) + + # 4. Output norm & projection + hidden_states = self.norm_out(hidden_states, temb) + hidden_states = self.proj_out(hidden_states) + + # 5. Unpatchify: [B, H'*W', C*p*p] -> [B, C, H, W] + hidden_states = hidden_states.reshape(batch_size, post_patch_height, post_patch_width, -1, p, p) + output = hidden_states.permute(0, 3, 1, 4, 2, 5).flatten(4, 5).flatten(2, 3) + + if not return_dict: + return (output,) + return Transformer2DModelOutput(sample=output) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """ + Load weights from pretrained checkpoint. + + This method handles the mapping from diffusers weight names to vllm-omni weight names, + especially for fused QKV projections. + """ + stacked_params_mapping = [ + # Fused QKV projection: to_q, to_k, to_v -> to_qkv + (".to_qkv", ".to_q", "q"), + (".to_qkv", ".to_k", "k"), + (".to_qkv", ".to_v", "v"), + ] + + params_dict = dict(self.named_parameters()) + + # Also include buffers (for any beta/eps parameters) + for name, buffer in self.named_buffers(): + params_dict[name] = buffer + + loaded_params: set[str] = set() + + for name, loaded_weight in weights: + # Handle fused QKV projections + for param_name, weight_name, shard_id in stacked_params_mapping: + if weight_name not in name: + continue + + # Map diffusers name to vllm-omni name + name = name.replace(weight_name, param_name) + + if name not in params_dict: + logger.warning(f"Skipping weight {name} - not found in model") + break + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight, shard_id) + break + else: + # Standard weight loading (not fused) + if name not in params_dict: + logger.warning(f"Skipping weight {name} - not found in model") + continue + + param = params_dict[name] + weight_loader = getattr(param, "weight_loader", default_weight_loader) + weight_loader(param, loaded_weight) + + loaded_params.add(name) + + return loaded_params + + def create_kv_cache(self) -> GlmImageKVCache: + """ + Create a KV cache for image editing. + + Returns a new GlmImageKVCache instance sized for this model's + number of transformer layers. Use this for image editing workflows. + + Example: + kv_cache = transformer.create_kv_cache() + kv_cache.set_mode("write") + transformer(condition_image, kv_cache=kv_cache) + kv_cache.set_mode("read") + for t in timesteps: + transformer(noisy_target, kv_cache=kv_cache) + kv_cache.clear() + + Returns: + GlmImageKVCache instance with correct number of layers. + """ + return GlmImageKVCache(num_layers=len(self.transformer_blocks)) + + @property + def num_layers(self) -> int: + """Return number of transformer layers.""" + return len(self.transformer_blocks) + + @property + def dtype(self) -> torch.dtype: + """Return dtype of model parameters.""" + return next(self.parameters()).dtype diff --git a/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py new file mode 100644 index 000000000..9a03a9349 --- /dev/null +++ b/vllm_omni/diffusion/models/glm_image/pipeline_glm_image.py @@ -0,0 +1,981 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +""" +GlmImagePipeline implementation for vLLM-Omni. + +This pipeline implements GLM-Image text-to-image generation with: +- AR stage: GlmImageForConditionalGeneration generates prior tokens +- DiT stage: GlmImageTransformer2DModel performs diffusion denoising +- VAE: AutoencoderKL decodes latents to images +""" + +from __future__ import annotations + +import inspect +import json +import logging +import os +import re +from collections.abc import Iterable + +import numpy as np +import PIL.Image +import torch +from diffusers.image_processor import VaeImageProcessor +from diffusers.models.autoencoders.autoencoder_kl import AutoencoderKL +from diffusers.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, +) +from diffusers.utils.torch_utils import randn_tensor +from torch import nn +from transformers import ( + ByT5Tokenizer, + GlmImageForConditionalGeneration, + GlmImageProcessor, + T5EncoderModel, +) + +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.distributed.parallel_state import ( + get_cfg_group, + get_classifier_free_guidance_rank, + get_classifier_free_guidance_world_size, +) +from vllm_omni.diffusion.distributed.utils import get_local_device +from vllm_omni.diffusion.model_loader.diffusers_loader import DiffusersPipelineLoader +from vllm_omni.diffusion.models.glm_image.glm_image_transformer import ( + GlmImageKVCache, + GlmImageTransformer2DModel, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest +from vllm_omni.model_executor.model_loader.weight_utils import ( + download_weights_from_hf_specific, +) + +logger = logging.getLogger(__name__) + + +def get_glm_image_pre_process_func(od_config: OmniDiffusionConfig): + """Get pre-processing function for GLM-Image pipeline. + + Pre-processes condition images before they are sent to the pipeline. + This is called by DiffusionEngine before batching requests. + """ + model_name = od_config.model + if os.path.exists(model_name): + model_path = model_name + else: + model_path = download_weights_from_hf_specific(model_name, None, ["*"]) + + vae_config_path = os.path.join(model_path, "vae/config.json") + with open(vae_config_path) as f: + vae_config = json.load(f) + block_out_channels = vae_config.get("block_out_channels", [128, 256, 512, 512]) + vae_scale_factor = 2 ** (len(block_out_channels) - 1) + + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) + # GLM-Image uses patch_size=2 for transformer + patch_size = 2 + + def pre_process_func(requests: list[OmniDiffusionRequest]): + """Pre-process condition images for Image Edit mode.""" + for req in requests: + images = req.pil_image + if images is None: + # Text-to-image mode, no preprocessing needed + continue + + if not isinstance(images, list): + images = [images] + + preprocessed = [] + height, width = None, None + + for img in images: + if isinstance(img, PIL.Image.Image): + img_h, img_w = img.size[::-1] # PIL is (width, height) + else: + img_h, img_w = img.shape[:2] + + # Align to multiple of vae_scale_factor * patch_size + multiple_of = vae_scale_factor * patch_size + img_h = (img_h // multiple_of) * multiple_of + img_w = (img_w // multiple_of) * multiple_of + + processed = image_processor.preprocess(img, height=img_h, width=img_w) + preprocessed.append(processed) + + # Use first image dimensions as default + if height is None: + height, width = img_h, img_w + + # Store in request + req.preprocessed_image = preprocessed + req.prompt_image = images # Keep original PIL images + if req.height is None: + req.height = height + if req.width is None: + req.width = width + + return requests + + return pre_process_func + + +def get_glm_image_post_process_func(od_config: OmniDiffusionConfig): + """Get post-processing function for GLM-Image pipeline.""" + model_name = od_config.model + if os.path.exists(model_name): + model_path = model_name + else: + model_path = download_weights_from_hf_specific(model_name, None, ["*"]) + + vae_config_path = os.path.join(model_path, "vae/config.json") + with open(vae_config_path) as f: + vae_config = json.load(f) + block_out_channels = vae_config.get("block_out_channels", [128, 256, 512, 512]) + vae_scale_factor = 2 ** (len(block_out_channels) - 1) + + image_processor = VaeImageProcessor(vae_scale_factor=vae_scale_factor) + + def post_process_func(images: torch.Tensor) -> list[PIL.Image.Image]: + return image_processor.postprocess(images, output_type="pil") + + return post_process_func + + +def calculate_shift( + image_seq_len: int, + base_seq_len: int = 256, + base_shift: float = 0.25, + max_shift: float = 0.75, +) -> float: + """Calculate timestep shift based on image sequence length.""" + m = (image_seq_len / base_seq_len) ** 0.5 + mu = m * max_shift + base_shift + return mu + + +def retrieve_timesteps( + scheduler, + num_inference_steps: int | None = None, + device: str | torch.device | None = None, + timesteps: list[int] | None = None, + sigmas: list[float] | None = None, + **kwargs, +) -> tuple[torch.Tensor, int]: + """ + Calls the scheduler's `set_timesteps` method and retrieves timesteps. + Handles custom timesteps and sigmas schedules. + """ + accepts_timesteps = "timesteps" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + accepts_sigmas = "sigmas" in set(inspect.signature(scheduler.set_timesteps).parameters.keys()) + + if timesteps is not None and sigmas is not None: + # Both provided - check if scheduler supports both + if not accepts_timesteps and not accepts_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep or sigma schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif timesteps is not None: + if not accepts_timesteps: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" timestep schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + elif sigmas is not None: + if not accepts_sigmas: + raise ValueError( + f"The current scheduler class {scheduler.__class__}'s `set_timesteps` does not support custom" + f" sigma schedules. Please check whether you are using the correct scheduler." + ) + scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs) + timesteps = scheduler.timesteps + num_inference_steps = len(timesteps) + else: + scheduler.set_timesteps(num_inference_steps, device=device, **kwargs) + timesteps = scheduler.timesteps + + return timesteps, num_inference_steps + + +def retrieve_latents( + encoder_output: torch.Tensor, + generator: torch.Generator | None = None, + sample_mode: str = "sample", +) -> torch.Tensor: + """Extract latents from VAE encoder output.""" + if hasattr(encoder_output, "latent_dist") and sample_mode == "sample": + return encoder_output.latent_dist.sample(generator) + elif hasattr(encoder_output, "latent_dist") and sample_mode == "argmax": + return encoder_output.latent_dist.mode() + elif hasattr(encoder_output, "latents"): + return encoder_output.latents + else: + raise AttributeError("Could not access latents of provided encoder_output") + + +class GlmImagePipeline(nn.Module): + """ + GLM-Image Pipeline for text-to-image and image-to-image generation. + + This pipeline integrates: + - AR model (GlmImageForConditionalGeneration): Generates prior image tokens + - Text encoder (T5EncoderModel): Encodes glyph/text embeddings + - DiT model (GlmImageTransformer2DModel): Diffusion transformer + - VAE (AutoencoderKL): Encodes/decodes images to/from latent space + + The pipeline flow: + 1. AR generates prior_token_ids from text prompt + 2. T5 encodes glyph text for text rendering + 3. DiT performs iterative denoising conditioned on prior tokens + 4. VAE decodes final latents to image + """ + + def __init__( + self, + *, + od_config: OmniDiffusionConfig, + prefix: str = "", + ): + super().__init__() + self.od_config = od_config + self.parallel_config = od_config.parallel_config + self.device = get_local_device() + + model = od_config.model + local_files_only = os.path.exists(model) + + if local_files_only: + model_path = model + else: + model_path = download_weights_from_hf_specific(model, od_config.revision, ["*"]) + + # Load scheduler + self.scheduler = FlowMatchEulerDiscreteScheduler.from_pretrained( + model_path, subfolder="scheduler", local_files_only=True + ) + + # Load AR model (vision_language_encoder) + logger.info("Loading GlmImageForConditionalGeneration (AR model)...") + self.vision_language_encoder = GlmImageForConditionalGeneration.from_pretrained( + model_path, + subfolder="vision_language_encoder", + local_files_only=True, + torch_dtype=torch.bfloat16, + ).to(self.device) + self.vision_language_encoder.eval() + + # Load processor for AR model + self.processor = GlmImageProcessor.from_pretrained(model_path, subfolder="processor", local_files_only=True) + + # Load text encoder (T5 for glyph embeddings) + logger.info("Loading T5EncoderModel (glyph encoder)...") + self.text_encoder = T5EncoderModel.from_pretrained( + model_path, + subfolder="text_encoder", + local_files_only=True, + torch_dtype=torch.bfloat16, + ).to(self.device) + self.text_encoder.eval() + + # Load tokenizer for glyph encoding + self.tokenizer = ByT5Tokenizer.from_pretrained(model_path, subfolder="tokenizer", local_files_only=True) + + # Load VAE + logger.info("Loading AutoencoderKL (VAE)...") + self.vae = AutoencoderKL.from_pretrained( + model_path, subfolder="vae", local_files_only=True, torch_dtype=torch.bfloat16 + ).to(self.device) + self.vae.eval() + + # Load transformer (DiT) + logger.info("Loading GlmImageTransformer2DModel (DiT)...") + self.transformer = GlmImageTransformer2DModel(od_config=od_config) + + # Weight sources for DiT loading + self.weights_sources = [ + DiffusersPipelineLoader.ComponentSource( + model_or_path=od_config.model, + subfolder="transformer", + revision=od_config.revision, + prefix="transformer.", + fall_back_to_pt=True, + ) + ] + + # Configure scale factors + self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1) + self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor) + self.default_sample_size = 128 + + # Get transformer config for patch size + self._patch_size = getattr(self.transformer, "patch_size", 2) + + # ==================== Input Validation ==================== + + def check_inputs( + self, + prompt: str | list[str] | None, + height: int | None, + width: int | None, + prompt_embeds: torch.Tensor | None = None, + ) -> None: + """Validate input arguments before generation.""" + # Check dimension alignment + multiple_of = self.vae_scale_factor * self._patch_size + if height is not None and height % multiple_of != 0: + logger.warning( + f"`height` should be divisible by {multiple_of} but is {height}. " + "Dimensions will be adjusted accordingly." + ) + if width is not None and width % multiple_of != 0: + logger.warning( + f"`width` should be divisible by {multiple_of} but is {width}. Dimensions will be adjusted accordingly." + ) + + # Check prompt/prompt_embeds mutual exclusivity + if prompt is not None and prompt_embeds is not None: + raise ValueError( + f"Cannot forward both `prompt`: {prompt} and `prompt_embeds`: {prompt_embeds}. " + "Please provide only one of the two." + ) + if prompt is None and prompt_embeds is None: + raise ValueError("Provide either `prompt` or `prompt_embeds`. Cannot leave both undefined.") + + # Check prompt type + if prompt is not None and not isinstance(prompt, (str, list)): + raise ValueError(f"`prompt` must be of type `str` or `list` but is {type(prompt)}") + + # ==================== AR Stage Methods ==================== + + @staticmethod + def _compute_generation_params( + image_grid_thw: torch.Tensor, + is_text_to_image: bool, + ) -> tuple[int, int, int, int]: + """ + Compute AR generation parameters from image grid. + + Args: + image_grid_thw: Image grid tensor of shape [N, 3] where each row is [t, h, w] + is_text_to_image: Whether this is text-to-image (vs image-to-image) + + Returns: + Tuple of (max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w) + """ + grid_sizes = [] + grid_hw = [] + + for i in range(image_grid_thw.shape[0]): + t, h, w = image_grid_thw[i].tolist() + grid_sizes.append(int(h * w)) + grid_hw.append((int(h), int(w))) + + if not is_text_to_image: + # Image-to-image: only generate target image tokens + max_new_tokens = grid_sizes[-1] + 1 + large_image_start_offset = 0 + target_grid_h, target_grid_w = grid_hw[-1] + else: + # Text-to-image: generate both small preview and large target + total_tokens = sum(grid_sizes) + max_new_tokens = total_tokens + 1 + large_image_start_offset = sum(grid_sizes[1:]) + target_grid_h, target_grid_w = grid_hw[0] + + return max_new_tokens, large_image_start_offset, target_grid_h, target_grid_w + + @staticmethod + def _extract_large_image_tokens( + outputs: torch.Tensor, input_length: int, large_image_start_offset: int, large_image_tokens: int + ) -> torch.Tensor: + """Extract large image tokens from AR output.""" + generated_tokens = outputs[0][input_length:] + large_image_start = large_image_start_offset + large_image_end = large_image_start + large_image_tokens + return generated_tokens[large_image_start:large_image_end] + + @staticmethod + def _upsample_token_ids(token_ids: torch.Tensor, token_h: int, token_w: int) -> torch.Tensor: + """Upsample token IDs by 2x using nearest neighbor interpolation.""" + token_ids = token_ids.view(1, 1, token_h, token_w) + token_ids = torch.nn.functional.interpolate(token_ids.float(), scale_factor=2, mode="nearest").to( + dtype=torch.long + ) + token_ids = token_ids.view(1, -1) + return token_ids + + @torch.inference_mode() + def generate_prior_tokens( + self, + prompt: str, + height: int, + width: int, + image: list[PIL.Image.Image] | None = None, + factor: int = 32, + ) -> tuple[torch.Tensor, list[torch.Tensor] | None]: + """ + Generate prior tokens using the AR model. + + Args: + prompt: Text prompt for generation + height: Target image height + width: Target image width + image: Optional condition images for image-to-image + factor: Token factor (default 32) + + Returns: + Tuple of (prior_token_ids, prior_token_image_ids) + prior_token_image_ids is a list of tensors, one per condition image + """ + device = self.vision_language_encoder.device + height = (height // factor) * factor + width = (width // factor) * factor + is_text_to_image = image is None or len(image) == 0 + + # Build message content + content = [] + if image is not None: + for img in image: + content.append({"type": "image", "image": img}) + content.append({"type": "text", "text": prompt}) + messages = [{"role": "user", "content": content}] + + # Apply chat template - processor will handle target dimensions and build grid + inputs = self.processor.apply_chat_template( + messages, + tokenize=True, + target_h=height, + target_w=width, + return_dict=True, + return_tensors="pt", + ).to(device) + + image_grid_thw = inputs.get("image_grid_thw") + + # Compute generation parameters from the full grid + max_new_tokens, large_image_offset, token_h, token_w = self._compute_generation_params( + image_grid_thw=image_grid_thw, is_text_to_image=is_text_to_image + ) + + # Process condition images if provided + # Use image_grid_thw[:-1] to exclude the target image grid (last entry) + prior_token_image_ids = None + if image is not None and image_grid_thw is not None and len(image_grid_thw) > 1: + # Get features only for condition images (exclude target image grid) + condition_grid = image_grid_thw[:-1] + prior_token_image_embed = self.vision_language_encoder.get_image_features( + inputs["pixel_values"], condition_grid + ) + prior_token_image_embed = torch.cat(prior_token_image_embed, dim=0) + flat_prior_token_image_ids = self.vision_language_encoder.get_image_tokens( + prior_token_image_embed, condition_grid + ) + # Split by image grid sizes and convert to list + split_sizes = (condition_grid.prod(dim=-1)).tolist() + prior_token_image_ids_list = torch.split(flat_prior_token_image_ids, split_sizes, dim=0) + # Convert to list with upsampling + prior_token_image_ids = [] + for i, token_ids in enumerate(prior_token_image_ids_list): + grid_t, grid_h, grid_w = condition_grid[i].tolist() + token_ids = token_ids.view(1, -1) + # Upsample 2x (from d32 to d64) + token_ids_upsampled = self._upsample_token_ids(token_ids, grid_h, grid_w) + prior_token_image_ids.append(token_ids_upsampled) + + # Generate with AR model + outputs = self.vision_language_encoder.generate( + **inputs, + max_new_tokens=max_new_tokens, + do_sample=True, + ) + + # Extract and upsample tokens + large_image_tokens = token_h * token_w + prior_token_ids_d32 = self._extract_large_image_tokens( + outputs, inputs["input_ids"].shape[-1], large_image_offset, large_image_tokens + ) + prior_token_ids = self._upsample_token_ids(prior_token_ids_d32, token_h, token_w) + + return prior_token_ids, prior_token_image_ids + + # ==================== Text Encoding Methods ==================== + + def get_glyph_texts(self, prompt: str | list[str]) -> list[str]: + """Extract text within quotes for glyph rendering.""" + prompt = prompt[0] if isinstance(prompt, list) else prompt + ocr_texts = ( + re.findall(r"'([^']*)'", prompt) + + re.findall(r"“([^“”]*)”", prompt) + + re.findall(r'"([^"]*)"', prompt) + + re.findall(r"「([^「」]*)」", prompt) + ) + return ocr_texts + + def _get_glyph_embeds( + self, + prompt: str | list[str], + max_sequence_length: int = 2048, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + ) -> torch.Tensor: + """Get glyph embeddings from T5 encoder for text rendering.""" + device = device or self.device + dtype = dtype or self.text_encoder.dtype + + glyph_texts = self.get_glyph_texts(prompt) + input_ids = self.tokenizer( + glyph_texts if len(glyph_texts) > 0 else [""], + max_length=max_sequence_length, + truncation=True, + ).input_ids + + # Pad to even length + input_ids = [[self.tokenizer.pad_token_id] * ((len(ids) + 1) % 2) + ids for ids in input_ids] + max_length = max(len(ids) for ids in input_ids) + + attention_mask = torch.tensor( + [[1] * len(ids) + [0] * (max_length - len(ids)) for ids in input_ids], + device=device, + ) + input_ids = torch.tensor( + [ids + [self.tokenizer.pad_token_id] * (max_length - len(ids)) for ids in input_ids], + device=device, + ) + + outputs = self.text_encoder(input_ids, attention_mask=attention_mask) + glyph_embeds = outputs.last_hidden_state[attention_mask.bool()].unsqueeze(0) + + return glyph_embeds.to(device=device, dtype=dtype) + + def encode_prompt( + self, + prompt: str | list[str], + do_classifier_free_guidance: bool = True, + num_images_per_prompt: int = 1, + prompt_embeds: torch.Tensor | None = None, + device: torch.device | None = None, + dtype: torch.dtype | None = None, + max_sequence_length: int = 2048, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + """Encode prompt into glyph embeddings for text rendering.""" + device = device or self.device + + prompt = [prompt] if isinstance(prompt, str) else prompt + batch_size = len(prompt) if prompt_embeds is None else prompt_embeds.shape[0] + + if prompt_embeds is None: + prompt_embeds = self._get_glyph_embeds(prompt, max_sequence_length, device, dtype) + + seq_len = prompt_embeds.size(1) + prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1) + prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + negative_prompt_embeds = None + if do_classifier_free_guidance: + negative_prompt = [""] * batch_size + negative_prompt_embeds = self._get_glyph_embeds(negative_prompt, max_sequence_length, device, dtype) + seq_len = negative_prompt_embeds.size(1) + negative_prompt_embeds = negative_prompt_embeds.repeat(1, num_images_per_prompt, 1) + negative_prompt_embeds = negative_prompt_embeds.view(batch_size * num_images_per_prompt, seq_len, -1) + + return prompt_embeds, negative_prompt_embeds + + # ==================== Latent Preparation ==================== + + def prepare_latents( + self, + batch_size: int, + num_channels_latents: int, + height: int, + width: int, + dtype: torch.dtype, + device: torch.device, + generator: torch.Generator | None, + latents: torch.Tensor | None = None, + ) -> torch.Tensor: + """Prepare random noise latents.""" + if latents is not None: + return latents.to(device) + + shape = ( + batch_size, + num_channels_latents, + int(height) // self.vae_scale_factor, + int(width) // self.vae_scale_factor, + ) + if isinstance(generator, list) and len(generator) != batch_size: + raise ValueError(f"Passed {len(generator)} generators but batch size is {batch_size}.") + latents = randn_tensor(shape, generator=generator, device=device, dtype=dtype) + return latents + + def diffuse( + self, + latents: torch.Tensor, + prior_token_id: torch.Tensor, + prompt_embeds: torch.Tensor, + negative_prompt_embeds: torch.Tensor | None, + timesteps: torch.Tensor, + target_size: torch.Tensor, + crop_coords: torch.Tensor, + guidance_scale: float, + do_classifier_free_guidance: bool, + kv_caches: GlmImageKVCache | None = None, + ) -> torch.Tensor: + """ + Denoising loop for diffusion process with CFG-Parallel support. + + Args: + latents: Initial noise latents + prior_token_id: Prior tokens generated by AR model + prompt_embeds: Encoded positive prompt embeddings (glyph embeddings) + negative_prompt_embeds: Encoded negative prompt embeddings + timesteps: Denoising timesteps + target_size: Target image size tensor [[height, width]] + crop_coords: Crop coordinates tensor + guidance_scale: CFG scale + do_classifier_free_guidance: Whether to apply CFG + kv_caches: Optional KV cache for Image Edit mode + + Returns: + Denoised latents ready for VAE decode + """ + # Prepare conditional/unconditional drop flags + prior_token_drop_cond = torch.full_like(prior_token_id, False, dtype=torch.bool) + prior_token_drop_uncond = torch.full_like(prior_token_id, True, dtype=torch.bool) + + transformer_dtype = self.transformer.dtype + + # Enable CFG-parallel: rank0 computes positive, rank1 computes negative + cfg_parallel_ready = do_classifier_free_guidance and get_classifier_free_guidance_world_size() > 1 + + for i, t in enumerate(timesteps): + latent_model_input = latents.to(transformer_dtype) + timestep = t.expand(latents.shape[0]) - 1 + + if cfg_parallel_ready: + cfg_group = get_cfg_group() + cfg_rank = get_classifier_free_guidance_rank() + + if cfg_rank == 0: + # Rank 0: Compute positive (conditional) prediction + local_pred = self.transformer( + hidden_states=latent_model_input, + encoder_hidden_states=prompt_embeds, + prior_token_id=prior_token_id, + prior_token_drop=prior_token_drop_cond, + timestep=timestep, + target_size=target_size, + crop_coords=crop_coords, + kv_cache=kv_caches, + return_dict=False, + )[0].float() + else: + # Rank 1: Compute negative (unconditional) prediction + local_pred = self.transformer( + hidden_states=latent_model_input, + encoder_hidden_states=negative_prompt_embeds, + prior_token_id=prior_token_id, + prior_token_drop=prior_token_drop_uncond, + timestep=timestep, + target_size=target_size, + crop_coords=crop_coords, + kv_cache=kv_caches, + return_dict=False, + )[0].float() + + # All-gather predictions from all ranks + gathered = cfg_group.all_gather(local_pred, separate_tensors=True) + + if cfg_rank == 0: + # Rank 0: Combine predictions and apply CFG + noise_pred_cond = gathered[0] + noise_pred_uncond = gathered[1] + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + # Scheduler step + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + # Broadcast updated latents to all ranks + cfg_group.broadcast(latents, src=0) + + else: + # Sequential CFG (single GPU or no CFG) + # Conditional forward pass + noise_pred_cond = self.transformer( + hidden_states=latent_model_input, + encoder_hidden_states=prompt_embeds, + prior_token_id=prior_token_id, + prior_token_drop=prior_token_drop_cond, + timestep=timestep, + target_size=target_size, + crop_coords=crop_coords, + kv_cache=kv_caches, + return_dict=False, + )[0].float() + + if do_classifier_free_guidance: + # Unconditional forward pass + noise_pred_uncond = self.transformer( + hidden_states=latent_model_input, + encoder_hidden_states=negative_prompt_embeds, + prior_token_id=prior_token_id, + prior_token_drop=prior_token_drop_uncond, + timestep=timestep, + target_size=target_size, + crop_coords=crop_coords, + kv_cache=kv_caches, + return_dict=False, + )[0].float() + + noise_pred = noise_pred_uncond + guidance_scale * (noise_pred_cond - noise_pred_uncond) + else: + noise_pred = noise_pred_cond + + # Scheduler step + latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0] + + return latents + + # ==================== Main Forward Pass ==================== + + def _prepare_condition_image_kv_cache( + self, + condition_images: list[torch.Tensor], + prior_token_image_ids: list[torch.Tensor], + prompt_embeds: torch.Tensor, + generator: torch.Generator | None = None, + ) -> GlmImageKVCache: + """ + Prepare KV cache by running condition images through transformer at timestep 0. + + This is used for Image Edit mode where we need to cache the condition image's + KV states for cross-attention during denoising. + + Args: + condition_images: List of preprocessed condition images + prior_token_image_ids: Prior token IDs for each condition image from AR model + prompt_embeds: Prompt embeddings (used to get dtype) + generator: Optional random generator + + Returns: + GlmImageKVCache with cached KV states from condition images + """ + kv_caches = self.transformer.create_kv_cache() + kv_caches.set_mode("write") + + # Prepare VAE normalization parameters + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.latent_channels, 1, 1) + .to(device=self.device, dtype=prompt_embeds.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.vae.config.latent_channels, 1, 1) + .to(device=self.device, dtype=prompt_embeds.dtype) + ) + + # Process each condition image through transformer to populate KV cache + for condition_image, condition_prior_token_id in zip(condition_images, prior_token_image_ids): + condition_image = condition_image.to(device=self.device, dtype=prompt_embeds.dtype) + + # Encode condition image to latent space + # Use argmax (mode) for deterministic encoding of condition images + condition_latent = retrieve_latents( + self.vae.encode(condition_image), generator=generator, sample_mode="argmax" + ) + condition_latent = (condition_latent - latents_mean) / latents_std + + # Run forward pass at timestep 0 to cache KV states + # Empty encoder_hidden_states since we only want to cache image features + _ = self.transformer( + hidden_states=condition_latent, + encoder_hidden_states=torch.zeros_like(prompt_embeds)[:1, :0, ...], + prior_token_id=condition_prior_token_id, + prior_token_drop=torch.full_like(condition_prior_token_id, False, dtype=torch.bool), + timestep=torch.zeros((1,), device=self.device), + target_size=torch.tensor([condition_image.shape[-2:]], device=self.device, dtype=prompt_embeds.dtype), + crop_coords=torch.zeros((1, 2), device=self.device, dtype=prompt_embeds.dtype), + kv_cache=kv_caches, + return_dict=False, + ) + + return kv_caches + + @torch.inference_mode() + def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + """ + Main generation forward pass. + + Args: + req: OmniDiffusionRequest with generation parameters + + Returns: + DiffusionOutput containing generated image + """ + prompt = req.prompt or "" + if isinstance(prompt, list): + prompt = prompt[0] if prompt else "" + + # Get pre-computed prompt embeddings if provided + prompt_embeds = req.prompt_embeds if isinstance(req.prompt_embeds, torch.Tensor) else None + + # Get condition images for Image Edit mode + # Use pre-processed images from pre_process_func + preprocessed_images = req.preprocessed_image + condition_images = getattr(req, "prompt_image", None) + img_height = req.height + img_width = req.width + + is_image_edit = preprocessed_images is not None + + # Use image dimensions as default if available + height = req.height or img_height or self.default_sample_size * self.vae_scale_factor + width = req.width or img_width or self.default_sample_size * self.vae_scale_factor + num_inference_steps = req.num_inference_steps or 50 + guidance_scale = req.guidance_scale or 1.5 + + # 0. Validate inputs + self.check_inputs(prompt=prompt, height=height, width=width, prompt_embeds=prompt_embeds) + + batch_size = 1 + do_classifier_free_guidance = guidance_scale > 1.0 + + # Set seed if provided + generator = None + if req.seed is not None: + generator = torch.Generator(device=self.device).manual_seed(req.seed) + + # 1. Get prior tokens - either from external source (multistage) or generate internally + # Check if prior_token_ids are provided externally (from AR stage in multistage mode) + external_prior_tokens = req.extra.get("prior_token_ids") if req.extra else None + external_prior_image_ids = req.extra.get("prior_token_image_ids") if req.extra else None + + if external_prior_tokens is not None: + # Multistage mode: use externally provided prior tokens from vLLM AR stage + logger.info("Using externally provided prior tokens from AR stage...") + prior_token_id = external_prior_tokens + if isinstance(prior_token_id, list): + prior_token_id = torch.tensor(prior_token_id, dtype=torch.long, device=self.device) + elif isinstance(prior_token_id, torch.Tensor): + prior_token_id = prior_token_id.to(device=self.device, dtype=torch.long) + # Ensure shape is [1, num_tokens] for batch processing + if prior_token_id.dim() == 1: + prior_token_id = prior_token_id.unsqueeze(0) + prior_token_image_ids = external_prior_image_ids + else: + # Single-stage mode: generate prior tokens with internal AR model + logger.info("Generating prior tokens with AR model...") + prior_token_id, prior_token_image_ids = self.generate_prior_tokens( + prompt=prompt, + image=condition_images, + height=height, + width=width, + ) + + # 2. Encode prompt for glyph embeddings + logger.info("Encoding prompt...") + prompt_embeds, negative_prompt_embeds = self.encode_prompt( + prompt, + do_classifier_free_guidance=do_classifier_free_guidance, + num_images_per_prompt=1, + prompt_embeds=prompt_embeds, + device=self.device, + dtype=self.transformer.dtype, + ) + + # 3. Prepare KV cache for Image Edit mode + kv_caches = None + if is_image_edit and prior_token_image_ids is not None: + logger.info("Preparing KV cache for Image Edit mode...") + kv_caches = self._prepare_condition_image_kv_cache( + condition_images=preprocessed_images, + prior_token_image_ids=prior_token_image_ids, + prompt_embeds=prompt_embeds, + generator=generator, + ) + # Switch to read mode for denoising + kv_caches.set_mode("read") + + # 4. Prepare latents + latent_channels = self.transformer.in_channels + latents = self.prepare_latents( + batch_size=batch_size, + num_channels_latents=latent_channels, + height=height, + width=width, + dtype=prompt_embeds.dtype, + device=self.device, + generator=generator, + ) + + # 5. Prepare timesteps + image_seq_len = ((height // self.vae_scale_factor) * (width // self.vae_scale_factor)) // (self._patch_size**2) + timesteps_array = np.linspace(self.scheduler.config.num_train_timesteps, 1.0, num_inference_steps + 1)[:-1] + timesteps_array = timesteps_array.astype(np.int64).astype(np.float32) + sigmas = timesteps_array / self.scheduler.config.num_train_timesteps + + mu = calculate_shift( + image_seq_len, + self.scheduler.config.get("base_image_seq_len", 256), + self.scheduler.config.get("base_shift", 0.25), + self.scheduler.config.get("max_shift", 0.75), + ) + timesteps, num_inference_steps = retrieve_timesteps( + self.scheduler, num_inference_steps, self.device, timesteps_array.tolist(), sigmas.tolist(), mu=mu + ) + + # 6. Prepare conditioning tensors + target_size = torch.tensor([[height, width]], dtype=prompt_embeds.dtype, device=self.device) + crop_coords = torch.zeros((1, 2), dtype=prompt_embeds.dtype, device=self.device) + + # 7. Denoising loop with CFG-parallel support + logger.info(f"Starting denoising loop with {num_inference_steps} steps...") + latents = self.diffuse( + latents=latents, + prior_token_id=prior_token_id, + prompt_embeds=prompt_embeds, + negative_prompt_embeds=negative_prompt_embeds, + timesteps=timesteps, + target_size=target_size, + crop_coords=crop_coords, + guidance_scale=guidance_scale, + do_classifier_free_guidance=do_classifier_free_guidance, + kv_caches=kv_caches, + ) + + # 8. VAE decode + logger.info("Decoding latents with VAE...") + latents = latents.to(self.vae.dtype) + latents_mean = ( + torch.tensor(self.vae.config.latents_mean) + .view(1, self.vae.config.latent_channels, 1, 1) + .to(latents.device, latents.dtype) + ) + latents_std = ( + torch.tensor(self.vae.config.latents_std) + .view(1, self.vae.config.latent_channels, 1, 1) + .to(latents.device, latents.dtype) + ) + latents = latents * latents_std + latents_mean + image = self.vae.decode(latents, return_dict=False, generator=generator)[0] + + # 9. Leave post-process to vllm-omni pipeline + + return DiffusionOutput(output=image) + + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]: + """Load transformer weights.""" + # Filter weights for transformer only + transformer_weights = ( + (name.replace("transformer.", ""), weight) for name, weight in weights if name.startswith("transformer.") + ) + return self.transformer.load_weights(transformer_weights) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index 5edd87a82..99b0ea4c7 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -29,6 +29,11 @@ _DIFFUSION_MODELS = { "pipeline_qwen_image_layered", "QwenImageLayeredPipeline", ), + "GlmImagePipeline": ( + "glm_image", + "pipeline_glm_image", + "GlmImagePipeline", + ), "ZImagePipeline": ( "z_image", "pipeline_z_image", @@ -117,6 +122,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "QwenImagePipeline": "get_qwen_image_post_process_func", "QwenImageEditPipeline": "get_qwen_image_edit_post_process_func", "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_post_process_func", + "GlmImagePipeline": "get_glm_image_post_process_func", "ZImagePipeline": "get_post_process_func", "OvisImagePipeline": "get_ovis_image_post_process_func", "WanPipeline": "get_wan22_post_process_func", @@ -133,6 +139,7 @@ _DIFFUSION_PRE_PROCESS_FUNCS = { # arch: pre_process_func # `pre_process_func` function must be placed in {mod_folder}/{mod_relname}.py, # where mod_folder and mod_relname are defined and mapped using `_DIFFUSION_MODELS` via the `arch` key + "GlmImagePipeline": "get_glm_image_pre_process_func", "QwenImageEditPipeline": "get_qwen_image_edit_pre_process_func", "QwenImageEditPlusPipeline": "get_qwen_image_edit_plus_pre_process_func", "LongCatImageEditPipeline": "get_longcat_image_edit_pre_process_func", diff --git a/vllm_omni/engine/__init__.py b/vllm_omni/engine/__init__.py index 79dc25cf4..0adfdb8d1 100644 --- a/vllm_omni/engine/__init__.py +++ b/vllm_omni/engine/__init__.py @@ -2,20 +2,14 @@ Engine components for vLLM-Omni. """ -import time -from collections.abc import Mapping from typing import Any import msgspec import torch from vllm.v1.engine import ( - EngineCoreEvent, + EngineCoreOutput, + EngineCoreOutputs, EngineCoreRequest, - FinishReason, - LogprobsLists, - LogprobsTensors, - SchedulerStats, - UtilityOutput, ) @@ -79,64 +73,9 @@ class OmniEngineCoreRequest(EngineCoreRequest): additional_information: AdditionalInformationPayload | None = None -class OmniEngineCoreOutput( - msgspec.Struct, - array_like=True, # type: ignore[call-arg] - omit_defaults=True, # type: ignore[call-arg] - gc=False, -): # type: ignore[call-arg] - request_id: str - new_token_ids: list[int] - - new_logprobs: LogprobsLists | None = None - new_prompt_logprobs_tensors: LogprobsTensors | None = None - +class OmniEngineCoreOutput(EngineCoreOutput): pooling_output: dict[str, torch.Tensor] | None = None - finish_reason: FinishReason | None = None - stop_reason: int | str | None = None - events: list[EngineCoreEvent] | None = None - kv_transfer_params: dict[str, Any] | None = None - trace_headers: Mapping[str, str] | None = None - # 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 - - -class OmniEngineCoreOutputs( - msgspec.Struct, - array_like=True, # type: ignore[call-arg] - omit_defaults=True, # type: ignore[call-arg] - gc=False, -): # type: ignore[call-arg] - # NOTE(Nick): We could consider ways to make this more compact, - # e.g. columnwise layout - - engine_index: int = 0 - - # [num_reqs] +class OmniEngineCoreOutputs(EngineCoreOutputs): outputs: list[OmniEngineCoreOutput] = [] - scheduler_stats: SchedulerStats | None = None - timestamp: float = 0.0 - - utility_output: UtilityOutput | None = None - finished_requests: set[str] | None = None - - # In DP case, used to signal that the current wave of requests - # has finished and the engines are paused. - wave_complete: int | None = None - # In DP case, used to signal that a request was received for an - # "old" wave, so the next wave needs to be started in other engines. - start_wave: int | None = None - - def __post_init__(self): - if self.timestamp == 0.0: - self.timestamp = time.monotonic() diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 6f22c6165..6aef7fb16 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -8,8 +8,6 @@ from vllm.outputs import PoolingRequestOutput from vllm.sampling_params import RequestOutputKind 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 from vllm.v1.engine.output_processor import OutputProcessor as VLLMOutputProcessor from vllm.v1.engine.output_processor import OutputProcessorOutput, RequestOutputCollector, RequestState from vllm.v1.engine.parallel_sampling import ParentRequest @@ -37,65 +35,6 @@ class OmniRequestState(RequestState): self.mm_type: str | None = None self.mm_accumulated: Dict[str, Any] | None = None - @classmethod - def from_new_request( - cls, - 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: - tokenizer = None - output_kind = sampling_params.output_kind - logprobs_processor = LogprobsProcessor.from_new_request( - tokenizer=tokenizer, - request=request, - ) - detokenizer = IncrementalDetokenizer.from_new_request( - tokenizer=tokenizer, - request=request, - ) - max_tokens_param = sampling_params.max_tokens - top_p = sampling_params.top_p - n = sampling_params.n - temperature = sampling_params.temperature - else: - logprobs_processor = None - detokenizer = None - max_tokens_param = None - top_p = None - n = None - temperature = None - assert request.pooling_params is not None - output_kind = request.pooling_params.output_kind - - return cls( - request_id=request.request_id, - parent_req=parent_req, - request_index=request_index, - lora_name=(request.lora_request.name if request.lora_request is not None else None), - output_kind=output_kind, - prompt=prompt, - prompt_token_ids=request.prompt_token_ids, - prompt_embeds=request.prompt_embeds, - logprobs_processor=logprobs_processor, - detokenizer=detokenizer, - max_tokens_param=max_tokens_param, - top_p=top_p, - n=n, - temperature=temperature, - 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: if payload is None: return diff --git a/vllm_omni/entrypoints/async_omni_llm.py b/vllm_omni/entrypoints/async_omni_llm.py index 567af0377..287f12b9e 100644 --- a/vllm_omni/entrypoints/async_omni_llm.py +++ b/vllm_omni/entrypoints/async_omni_llm.py @@ -10,7 +10,7 @@ import vllm.envs as envs from vllm.config import VllmConfig from vllm.logger import init_logger from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalRegistry -from vllm.tokenizers import init_tokenizer_from_config +from vllm.tokenizers import cached_tokenizer_from_config from vllm.tracing import init_tracer from vllm.transformers_utils.config import maybe_register_config_serialize_by_value from vllm.usage.usage_lib import UsageContext @@ -111,7 +111,7 @@ class AsyncOmniLLM(AsyncLLM): tokenizer = None else: # Tokenizer (+ ensure liveness if running in another process). - tokenizer = init_tokenizer_from_config(model_config=vllm_config.model_config) + tokenizer = cached_tokenizer_from_config(model_config=vllm_config.model_config) # InputProcessor (converts Inputs --> EngineCoreRequests). self.input_processor = OmniInputProcessor( diff --git a/vllm_omni/entrypoints/chat_utils.py b/vllm_omni/entrypoints/chat_utils.py index 6517ae666..0fdef5edb 100644 --- a/vllm_omni/entrypoints/chat_utils.py +++ b/vllm_omni/entrypoints/chat_utils.py @@ -22,7 +22,6 @@ from vllm.entrypoints.chat_utils import ( _postprocess_messages, _ToolParser, ) -from vllm.transformers_utils.tokenizer import AnyTokenizer class OmniAsyncMultiModalItemTracker(AsyncMultiModalItemTracker): @@ -129,7 +128,6 @@ class OmniAsyncMultiModalContentParser(AsyncMultiModalContentParser): def parse_chat_messages_futures( messages: list[ChatCompletionMessageParam], model_config: ModelConfig, - tokenizer: AnyTokenizer, content_format: _ChatTemplateContentFormat, mm_processor_kwargs: dict[str, Any] | None = None, ) -> tuple[ @@ -138,7 +136,7 @@ def parse_chat_messages_futures( MultiModalUUIDDict | None, ]: conversation: list[ConversationMessage] = [] - mm_tracker = OmniAsyncMultiModalItemTracker(model_config, tokenizer) + mm_tracker = OmniAsyncMultiModalItemTracker(model_config) for msg in messages: sub_messages = _parse_chat_message_content( diff --git a/vllm_omni/entrypoints/omni_llm.py b/vllm_omni/entrypoints/omni_llm.py index 05a48feee..8db4b11d9 100644 --- a/vllm_omni/entrypoints/omni_llm.py +++ b/vllm_omni/entrypoints/omni_llm.py @@ -1,12 +1,15 @@ +from collections.abc import Callable from typing import Any import cloudpickle from pydantic import ValidationError +from tqdm import tqdm # External library imports (vLLM) from vllm.config import CompilationConfig, StructuredOutputsConfig, is_init_field from vllm.entrypoints.llm import LLM from vllm.logger import init_logger +from vllm.outputs import PoolingRequestOutput, RequestOutput from vllm.plugins.io_processors import get_io_processor from vllm.usage.usage_lib import UsageContext from vllm.utils.counter import Counter @@ -190,3 +193,47 @@ class OmniLLM(LLM): self.close() except Exception as e: logger.debug("[Orchestrator] __del__ close() raised: %s", e, exc_info=True) + + def _run_engine(self, *, use_tqdm: bool | Callable[..., tqdm] = True) -> list[RequestOutput | PoolingRequestOutput]: + # Initialize tqdm. + if use_tqdm: + num_requests = self.llm_engine.get_num_unfinished_requests() + tqdm_func = use_tqdm if callable(use_tqdm) else tqdm + pbar = tqdm_func( + total=num_requests, + desc="Processed prompts", + dynamic_ncols=True, + postfix=(f"est. speed input: {0:.2f} toks/s, output: {0:.2f} toks/s"), + ) + + # Run the engine. + outputs: list[RequestOutput | PoolingRequestOutput] = [] + total_in_toks = 0 + total_out_toks = 0 + while self.llm_engine.has_unfinished_requests(): + step_outputs = self.llm_engine.step() + for output in step_outputs: + if output.finished: + outputs.append(output) + if use_tqdm: + if isinstance(output, RequestOutput): + # Calculate tokens only for RequestOutput + n = len(output.outputs) + assert output.prompt_token_ids is not None + total_in_toks += len(output.prompt_token_ids) * n + in_spd = total_in_toks / pbar.format_dict["elapsed"] + total_out_toks += sum(len(stp.token_ids) for stp in output.outputs) + out_spd = total_out_toks / pbar.format_dict["elapsed"] + pbar.postfix = f"est. speed input: {in_spd:.2f} toks/s, output: {out_spd:.2f} toks/s" + pbar.update(n) + else: + pbar.update(1) + if pbar.n == num_requests: + pbar.refresh() + + if use_tqdm: + pbar.close() + # Sort the outputs by the int part of request ID which is in format of 'int-uuid'. + # This is necessary because some requests may be finished earlier than + # its previous requests. + return sorted(outputs, key=lambda x: int(x.request_id.split("-")[0])) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 49958eb66..888d39085 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -1,3 +1,5 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project import multiprocessing import multiprocessing.forkserver as forkserver import os @@ -14,31 +16,48 @@ import vllm.envs as envs from fastapi import Depends, HTTPException, Request from fastapi.responses import JSONResponse, StreamingResponse from starlette.datastructures import State -from vllm.config import VllmConfig +from starlette.routing import Route from vllm.engine.protocol import EngineClient -from vllm.entrypoints.chat_utils import load_chat_template, resolve_hf_chat_template, resolve_mistral_chat_template +from vllm.entrypoints.anthropic.serving_messages import AnthropicServingMessages from vllm.entrypoints.launcher import serve_http from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.openai.api_server import ( base, build_app, load_log_config, - maybe_register_tokenizer_info_endpoint, router, setup_server, - validate_json_request, ) -from vllm.entrypoints.openai.protocol import ChatCompletionRequest, ChatCompletionResponse, ErrorResponse -from vllm.entrypoints.openai.serving_models import BaseModelPath, LoRAModulePath, OpenAIServingModels -from vllm.entrypoints.openai.tool_parsers import ToolParserManager - -# yapf conflicts with isort for this block -# yapf: disable -# yapf: enable +from vllm.entrypoints.openai.orca_metrics import metrics_header +from vllm.entrypoints.openai.protocol import ( + ChatCompletionRequest, + ChatCompletionResponse, + ErrorResponse, +) +from vllm.entrypoints.openai.serving_completion import OpenAIServingCompletion +from vllm.entrypoints.openai.serving_models import BaseModelPath, OpenAIServingModels +from vllm.entrypoints.openai.serving_responses import OpenAIServingResponses +from vllm.entrypoints.openai.serving_transcription import ( + OpenAIServingTranscription, + OpenAIServingTranslation, +) +from vllm.entrypoints.openai.utils import validate_json_request +from vllm.entrypoints.pooling.classify.serving import ServingClassification +from vllm.entrypoints.pooling.embed.serving import OpenAIServingEmbedding +from vllm.entrypoints.pooling.pooling.serving import OpenAIServingPooling +from vllm.entrypoints.pooling.score.serving import ServingScores +from vllm.entrypoints.serve.disagg.serving import ServingTokens +from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization from vllm.entrypoints.tool_server import DemoToolServer, MCPToolServer, ToolServer -from vllm.entrypoints.utils import load_aware_call, with_cancellation +from vllm.entrypoints.utils import ( + load_aware_call, + process_chat_template, + process_lora_modules, + with_cancellation, +) from vllm.logger import init_logger -from vllm.tokenizers import MistralTokenizer +from vllm.tasks import POOLING_TASKS +from vllm.tool_parsers import ToolParserManager from vllm.utils.system_utils import decorate_logs from vllm_omni.entrypoints.async_omni import AsyncOmni @@ -57,6 +76,29 @@ from vllm_omni.entrypoints.openai.serving_speech import OmniOpenAIServingSpeech logger = init_logger(__name__) +ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL = "endpoint-load-metrics-format" + + +def _remove_route_from_router(router_obj, path: str, methods: set[str] | None = None): + """Remove a route from the router by path and optionally by methods. + + This is needed because vllm's api_server registers routes when imported, + and we need to override some routes (like /v1/chat/completions) with + omni-specific implementations. + """ + routes_to_remove = [] + for route in router_obj.routes: + if isinstance(route, Route) and route.path == path: + if methods is None or (hasattr(route, "methods") and route.methods & methods): + routes_to_remove.append(route) + + for route in routes_to_remove: + router_obj.routes.remove(route) + + +# Remove vllm's /v1/chat/completions route so we can register our own omni version +_remove_route_from_router(router, "/v1/chat/completions", {"POST"}) + # Server entry points @@ -88,6 +130,10 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, if args.tool_parser_plugin and len(args.tool_parser_plugin) > 3: ToolParserManager.import_tool_parser(args.tool_parser_plugin) + if args.reasoning_parser_plugin and len(args.reasoning_parser_plugin) > 3: + from vllm.reasoning import ReasoningParserManager + + ReasoningParserManager.import_reasoning_parser(args.reasoning_parser_plugin) # Load logging config for uvicorn if specified log_config = load_log_config(args.log_config_file) @@ -98,11 +144,11 @@ async def omni_run_server_worker(listen_address, sock, args, client_config=None, args, client_config=client_config, ) as engine_client: - maybe_register_tokenizer_info_endpoint(args) app = build_app(args) + await omni_init_app_state(engine_client, app.state, args) + vllm_config = await engine_client.get_vllm_config() - await omni_init_app_state(engine_client, vllm_config, app.state, args) # Check if pure diffusion mode (vllm_config will be None) is_pure_diffusion = vllm_config is None @@ -233,7 +279,6 @@ async def build_async_omni_from_stage_config( async def omni_init_app_state( engine_client: EngineClient, - vllm_config: VllmConfig | None, state: State, args: Namespace, ) -> None: @@ -246,10 +291,12 @@ async def omni_init_app_state( Args: engine_client: Engine client instance (AsyncOmni) - vllm_config: vLLM configuration object (may be None for pure diffusion) state: FastAPI application state object to initialize args: Parsed command-line arguments """ + # Get vllm_config from engine_client (following 0.14.0 pattern) + vllm_config = await engine_client.get_vllm_config() + # Detect if it's pure Diffusion mode (single stage and is Diffusion) is_pure_diffusion = False if hasattr(engine_client, "stage_configs") and engine_client.stage_configs: @@ -273,6 +320,7 @@ async def omni_init_app_state( base_model_paths = [BaseModelPath(name=name, model_path=args.model) for name in served_model_names] state.engine_client = engine_client state.log_stats = not args.disable_log_stats + state.args = args # For omni models state.stage_configs = engine_client.stage_configs if hasattr(engine_client, "stage_configs") else None @@ -302,32 +350,18 @@ async def omni_init_app_state( logger.warning("vllm_config is None, some features may not work correctly") state.vllm_config = vllm_config - if vllm_config is not None: - _model_config = vllm_config.model_config - resolved_chat_template = load_chat_template(args.chat_template) - if resolved_chat_template is not None and vllm_config is not None: - # Get the tokenizer to check official template - tokenizer = await engine_client.get_tokenizer() + # Get supported tasks + supported_tasks: set[str] = {"generate"} + if hasattr(engine_client, "get_supported_tasks"): + supported_tasks = set(await engine_client.get_supported_tasks()) + logger.info("Supported tasks: %s", supported_tasks) - if tokenizer is not None: - if isinstance(tokenizer, MistralTokenizer): - # The warning is logged in resolve_mistral_chat_template. - resolved_chat_template = resolve_mistral_chat_template(chat_template=resolved_chat_template) - else: - hf_chat_template = resolve_hf_chat_template( - tokenizer=tokenizer, - chat_template=None, - tools=None, - model_config=vllm_config.model_config, - ) - - if hf_chat_template != resolved_chat_template: - logger.warning( - "Using supplied chat template: %s\nIt is different from official chat template '%s'. This discrepancy may lead to performance degradation.", # noqa: E501 - resolved_chat_template, - args.model, - ) + resolved_chat_template = await process_chat_template( + args.chat_template, + engine_client, + vllm_config.model_config if vllm_config is not None else None, + ) if args.tool_server == "demo": tool_server: ToolServer | None = DemoToolServer() @@ -340,23 +374,12 @@ async def omni_init_app_state( tool_server = None # Merge default_mm_loras into the static lora_modules - default_mm_loras = {} - if vllm_config is not None and vllm_config.lora_config is not None: - default_mm_loras = vllm_config.lora_config.default_mm_loras - - lora_modules = args.lora_modules - if default_mm_loras: - default_mm_lora_paths = [ - LoRAModulePath( - name=modality, - path=lora_path, - ) - for modality, lora_path in default_mm_loras.items() - ] - if args.lora_modules is None: - lora_modules = default_mm_lora_paths - else: - lora_modules += default_mm_lora_paths + default_mm_loras = ( + vllm_config.lora_config.default_mm_loras + if vllm_config is not None and vllm_config.lora_config is not None + else {} + ) + lora_modules = process_lora_modules(args.lora_modules, default_mm_loras) # Ensure input_processor, io_processor, and model_config exist for OpenAIServingModels compatibility if ( @@ -415,24 +438,182 @@ async def omni_init_app_state( lora_modules=lora_modules, ) await state.openai_serving_models.init_static_loras() - state.openai_serving_chat = OmniOpenAIServingChat( + + state.openai_serving_responses = ( + OpenAIServingResponses( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_auto_tools=args.enable_auto_tool_choice, + tool_parser=args.tool_call_parser, + tool_server=tool_server, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_force_include_usage=args.enable_force_include_usage, + enable_log_outputs=args.enable_log_outputs, + log_error_stack=args.log_error_stack, + ) + if "generate" in supported_tasks + else None + ) + state.openai_serving_chat = ( + OmniOpenAIServingChat( + engine_client, + state.openai_serving_models, + args.response_role, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + default_chat_template_kwargs=args.default_chat_template_kwargs, + trust_request_chat_template=args.trust_request_chat_template, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_auto_tools=args.enable_auto_tool_choice, + exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_force_include_usage=args.enable_force_include_usage, + enable_log_outputs=args.enable_log_outputs, + enable_log_deltas=args.enable_log_deltas, + log_error_stack=args.log_error_stack, + ) + if "generate" in supported_tasks + else None + ) + # Warm up chat template processing to avoid first-request latency + if state.openai_serving_chat is not None: + await state.openai_serving_chat.warmup() + + state.openai_serving_completion = ( + OpenAIServingCompletion( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_force_include_usage=args.enable_force_include_usage, + log_error_stack=args.log_error_stack, + ) + if "generate" in supported_tasks + else None + ) + state.openai_serving_pooling = ( + OpenAIServingPooling( + engine_client, + state.openai_serving_models, + supported_tasks=supported_tasks, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + log_error_stack=args.log_error_stack, + ) + if any(task in POOLING_TASKS for task in supported_tasks) + else None + ) + state.openai_serving_embedding = ( + OpenAIServingEmbedding( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + log_error_stack=args.log_error_stack, + ) + if "embed" in supported_tasks + else None + ) + state.openai_serving_classification = ( + ServingClassification( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + log_error_stack=args.log_error_stack, + ) + if "classify" in supported_tasks + else None + ) + state.openai_serving_scores = ( + ServingScores( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + score_template=resolved_chat_template, + log_error_stack=args.log_error_stack, + ) + if ("embed" in supported_tasks or "score" in supported_tasks) + else None + ) + state.openai_serving_tokenization = OpenAIServingTokenization( engine_client, state.openai_serving_models, - args.response_role, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, trust_request_chat_template=args.trust_request_chat_template, - return_tokens_as_token_ids=args.return_tokens_as_token_ids, - enable_auto_tools=args.enable_auto_tool_choice, - exclude_tools_when_tool_choice_none=args.exclude_tools_when_tool_choice_none, - tool_parser=args.tool_call_parser, - reasoning_parser=args.structured_outputs_config.reasoning_parser, - enable_prompt_tokens_details=args.enable_prompt_tokens_details, - enable_force_include_usage=args.enable_force_include_usage, - enable_log_outputs=args.enable_log_outputs, log_error_stack=args.log_error_stack, ) + state.openai_serving_transcription = ( + OpenAIServingTranscription( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + log_error_stack=args.log_error_stack, + enable_force_include_usage=args.enable_force_include_usage, + ) + if "transcription" in supported_tasks + else None + ) + state.openai_serving_translation = ( + OpenAIServingTranslation( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + log_error_stack=args.log_error_stack, + enable_force_include_usage=args.enable_force_include_usage, + ) + if "transcription" in supported_tasks + else None + ) + state.anthropic_serving_messages = ( + AnthropicServingMessages( + engine_client, + state.openai_serving_models, + args.response_role, + request_logger=request_logger, + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + enable_auto_tools=args.enable_auto_tool_choice, + tool_parser=args.tool_call_parser, + reasoning_parser=args.structured_outputs_config.reasoning_parser, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_force_include_usage=args.enable_force_include_usage, + ) + if "generate" in supported_tasks + else None + ) + state.serving_tokens = ( + ServingTokens( + engine_client, + state.openai_serving_models, + request_logger=request_logger, + return_tokens_as_token_ids=args.return_tokens_as_token_ids, + log_error_stack=args.log_error_stack, + enable_prompt_tokens_details=args.enable_prompt_tokens_details, + enable_log_outputs=args.enable_log_outputs, + force_no_detokenize=args.tokens_only, + ) + if "generate" in supported_tasks + else None + ) state.openai_serving_speech = OmniOpenAIServingSpeech( engine_client, state.openai_serving_models, request_logger=request_logger @@ -463,6 +644,7 @@ def Omnispeech(request: Request) -> OmniOpenAIServingSpeech | None: @with_cancellation @load_aware_call async def create_chat_completion(request: ChatCompletionRequest, raw_request: Request): + metrics_header_format = raw_request.headers.get(ENDPOINT_LOAD_METRICS_FORMAT_HEADER_LABEL, "") handler = Omnichat(raw_request) if handler is None: return base(raw_request).create_error_response(message="The model does not support Chat Completions API") @@ -474,7 +656,8 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re if isinstance(generator, ErrorResponse): return JSONResponse( - content=generator.model_dump(), status_code=generator.code if hasattr(generator, "code") else 400 + content=generator.model_dump(), + status_code=generator.error.code if generator.error else 400, ) elif isinstance(generator, ChatCompletionResponse): @@ -490,18 +673,27 @@ async def create_chat_completion(request: ChatCompletionRequest, raw_request: Re try: # Use serialize_as_any=True to bypass type checking response_dict = generator.model_dump(mode="json", serialize_as_any=True, warnings="none") - return JSONResponse(content=response_dict) + return JSONResponse( + content=response_dict, + headers=metrics_header(metrics_header_format), + ) except Exception: # Fallback: convert to JSON string and parse back to avoid any serialization issues try: response_json = generator.model_dump_json(warnings="none", serialize_as_any=True) response_dict = json_lib.loads(response_json) - return JSONResponse(content=response_dict) + return JSONResponse( + content=response_dict, + headers=metrics_header(metrics_header_format), + ) except Exception: # Last resort: regular dump with warnings suppressed with warnings_module.catch_warnings(): warnings_module.filterwarnings("ignore", category=UserWarning) - return JSONResponse(content=generator.model_dump(mode="json", warnings="none")) + return JSONResponse( + content=generator.model_dump(mode="json", warnings="none"), + headers=metrics_header(metrics_header_format), + ) return StreamingResponse(content=generator, media_type="text/event-stream") diff --git a/vllm_omni/entrypoints/openai/protocol/chat_completion.py b/vllm_omni/entrypoints/openai/protocol/chat_completion.py index 9f6076249..d0c83f56f 100644 --- a/vllm_omni/entrypoints/openai/protocol/chat_completion.py +++ b/vllm_omni/entrypoints/openai/protocol/chat_completion.py @@ -1,4 +1,6 @@ -from vllm.entrypoints.openai.protocol import ChatCompletionStreamResponse +from vllm.entrypoints.openai.protocol import ( + ChatCompletionStreamResponse, +) class OmniChatCompletionStreamResponse(ChatCompletionStreamResponse): diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 6fb9750cc..6320cb73a 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -29,7 +29,10 @@ from vllm.entrypoints.chat_utils import ( make_tool_call_id, resolve_chat_template_content_format, ) -from vllm.entrypoints.harmony_utils import get_streamable_parser_for_assistant, parse_chat_output +from vllm.entrypoints.openai.parser.harmony_utils import ( + get_streamable_parser_for_assistant, + parse_chat_output, +) from vllm.entrypoints.openai.protocol import ( ChatCompletionNamedToolChoiceParam, ChatCompletionRequest, @@ -41,29 +44,24 @@ from vllm.entrypoints.openai.protocol import ( DeltaFunctionCall, DeltaMessage, DeltaToolCall, + ErrorInfo, ErrorResponse, FunctionCall, FunctionDefinition, PromptTokenUsageInfo, RequestResponseMetadata, + ResponsesRequest, ToolCall, UsageInfo, ) from vllm.entrypoints.openai.serving_chat import OpenAIServingChat from vllm.entrypoints.openai.serving_engine import ( ChatLikeRequest, - EngineTokensPrompt, - RequestPrompt, - ResponsesRequest, - TextTokensPrompt, clamp_prompt_logprobs, - is_list_of, ) -from vllm.entrypoints.openai.tool_parsers import ToolParser -from vllm.entrypoints.openai.tool_parsers.mistral_tool_parser import MistralToolCall from vllm.entrypoints.openai.utils import maybe_filter_parallel_tool_calls from vllm.entrypoints.utils import should_include_usage -from vllm.inputs.data import PromptType +from vllm.inputs.data import PromptType, TokensPrompt from vllm.logger import init_logger from vllm.lora.request import LoRARequest from vllm.outputs import RequestOutput @@ -75,8 +73,10 @@ from vllm.tokenizers.mistral import ( truncate_tool_call_ids, validate_request_params, ) +from vllm.tool_parsers import ToolParser +from vllm.tool_parsers.mistral_tool_parser import MistralToolCall from vllm.transformers_utils.tokenizer import AnyTokenizer -from vllm.utils.collection_utils import as_list +from vllm.utils.collection_utils import as_list, is_list_of from vllm_omni.entrypoints.chat_utils import parse_chat_messages_futures from vllm_omni.entrypoints.openai.audio_utils_mixin import AudioMixin @@ -177,52 +177,68 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): truncate_tool_call_ids(request) validate_request_params(request) - if ( - request.tool_choice == "auto" - and not (self.enable_auto_tools and tool_parser is not None) - and not isinstance(tokenizer, MistralTokenizer) - and not self.use_harmony + # Check if tool parsing is unavailable (common condition) + tool_parsing_unavailable = ( + tool_parser is None and not isinstance(tokenizer, MistralTokenizer) and not self.use_harmony + ) + + # Validate tool_choice when tool parsing is required but unavailable + if tool_parsing_unavailable and request.tool_choice not in ( + None, + "none", ): - # for hf tokenizers, "auto" tools requires - # --enable-auto-tool-choice and --tool-call-parser - return self.create_error_response( - '"auto" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set' - ) + if request.tool_choice == "auto" and not self.enable_auto_tools: + # for hf tokenizers, "auto" tools requires + # --enable-auto-tool-choice and --tool-call-parser + return self.create_error_response( + '"auto" tool choice requires --enable-auto-tool-choice and --tool-call-parser to be set' + ) + elif request.tool_choice != "auto": + # "required" or named tool requires tool parser + return self.create_error_response( + f'tool_choice="{request.tool_choice}" requires --tool-call-parser to be set' + ) if request.tools is None or (request.tool_choice == "none" and self.exclude_tools_when_tool_choice_none): tool_dicts = None else: tool_dicts = [tool.model_dump() for tool in request.tools] - # Common case. - request_chat_template = request.chat_template - chat_template_kwargs = request.chat_template_kwargs - if not self.trust_request_chat_template and ( - request_chat_template is not None - or (chat_template_kwargs and chat_template_kwargs.get("chat_template") is not None) - ): - return self.create_error_response( - "Chat template is passed with request, but --trust-request-chat-template is not set. " - "Refused request with untrusted chat template." + if not self.use_harmony: + error_check_ret = self._validate_chat_template( + request_chat_template=request.chat_template, + chat_template_kwargs=request.chat_template_kwargs, + trust_request_chat_template=self.trust_request_chat_template, ) - ( - conversation, - request_prompts, - engine_prompts, - ) = await self._preprocess_chat( - request, - tokenizer, - request.messages, - chat_template=request_chat_template or self.chat_template, - chat_template_content_format=self.chat_template_content_format, - add_generation_prompt=request.add_generation_prompt, - continue_final_message=request.continue_final_message, - tool_dicts=tool_dicts, - documents=request.documents, - chat_template_kwargs=request.chat_template_kwargs, - tool_parser=tool_parser, - add_special_tokens=request.add_special_tokens, - ) + if error_check_ret is not None: + return error_check_ret + + chat_template_kwargs = request.chat_template_kwargs or {} + chat_template_kwargs.update(reasoning_effort=request.reasoning_effort) + + ( + conversation, + request_prompts, + engine_prompts, + ) = await self._preprocess_chat( + request, + tokenizer, + request.messages, + chat_template=request.chat_template or self.chat_template, + chat_template_content_format=self.chat_template_content_format, + add_generation_prompt=request.add_generation_prompt, + continue_final_message=request.continue_final_message, + tool_dicts=tool_dicts, + documents=getattr(request, "documents", None), + chat_template_kwargs=chat_template_kwargs, + default_chat_template_kwargs=self.default_chat_template_kwargs, + tool_parser=tool_parser, + add_special_tokens=request.add_special_tokens, + ) + else: + should_include_tools = tool_dicts is not None + conversation, engine_prompts = self._make_request_with_harmony(request, should_include_tools) + request_prompts = [engine_prompt.get("prompt_token_ids", []) for engine_prompt in engine_prompts] except (ValueError, TypeError, RuntimeError, jinja2.TemplateError) as e: logger.exception("Error in preprocessing prompt inputs") @@ -314,12 +330,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tool_dicts: list[dict[str, Any]] | None = None, documents: list[dict[str, str]] | None = None, chat_template_kwargs: dict[str, Any] | None = None, + default_chat_template_kwargs: dict[str, Any] | None = None, tool_parser: Callable[[TokenizerLike], ToolParser] | None = None, add_special_tokens: bool = False, ) -> tuple[ list[ConversationMessage], - Sequence[RequestPrompt], - list[EngineTokensPrompt], + Sequence[PromptType], + list[TokensPrompt], ]: model_config = self.model_config @@ -333,11 +350,17 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): conversation, mm_data_future, mm_uuids = parse_chat_messages_futures( messages, model_config, - tokenizer, content_format=resolved_content_format, mm_processor_kwargs=getattr(request, "mm_processor_kwargs", None), ) + # Merge default_chat_template_kwargs with request-provided kwargs + # Request kwargs take precedence over defaults + merged_kwargs = self._prepare_extra_chat_template_kwargs( + chat_template_kwargs, + default_chat_template_kwargs, + ) + _chat_template_kwargs: dict[str, Any] = dict( chat_template=chat_template, add_generation_prompt=add_generation_prompt, @@ -345,7 +368,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tools=tool_dicts, documents=documents, ) - _chat_template_kwargs.update(chat_template_kwargs or {}) + _chat_template_kwargs.update(merged_kwargs) request_prompt: str | list[int] @@ -388,7 +411,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "Prompt has to be a string", "when the tokenizer is not initialised", ) - prompt_inputs = TextTokensPrompt(prompt=request_prompt, prompt_token_ids=[1]) + prompt_inputs = TokensPrompt(prompt=request_prompt, prompt_token_ids=[1]) elif isinstance(request_prompt, str): prompt_inputs = await self._tokenize_prompt_input_async( request, @@ -399,20 +422,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): else: # For MistralTokenizer assert is_list_of(request_prompt, int), "Prompt has to be either a string or a list of token ids" - prompt_inputs = TextTokensPrompt( + prompt_inputs = TokensPrompt( prompt=tokenizer.decode(request_prompt), prompt_token_ids=request_prompt, ) - engine_prompt = EngineTokensPrompt(prompt_token_ids=prompt_inputs["prompt_token_ids"]) + engine_prompt = TokensPrompt(prompt_token_ids=prompt_inputs["prompt_token_ids"]) if mm_data is not None: engine_prompt["multi_modal_data"] = mm_data if mm_uuids is not None: engine_prompt["multi_modal_uuids"] = mm_uuids - if request.mm_processor_kwargs is not None: - engine_prompt["mm_processor_kwargs"] = request.mm_processor_kwargs + mm_processor_kwargs = getattr(request, "mm_processor_kwargs", None) + if mm_processor_kwargs is not None: + engine_prompt["mm_processor_kwargs"] = mm_processor_kwargs if hasattr(request, "cache_salt") and request.cache_salt is not None: engine_prompt["cache_salt"] = request.cache_salt @@ -513,7 +537,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): def _log_inputs( self, request_id: str, - inputs: RequestPrompt | PromptType, + inputs: PromptType, params_list: list[SamplingParams] | None, lora_request: LoRARequest | None, ) -> None: @@ -599,9 +623,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): try: if self.reasoning_parser: + chat_template_kwargs = self._prepare_extra_chat_template_kwargs( + request.chat_template_kwargs, + self.default_chat_template_kwargs, + ) reasoning_parser = self.reasoning_parser( tokenizer, - chat_template_kwargs=request.chat_template_kwargs, # type: ignore + chat_template_kwargs=chat_template_kwargs, # type: ignore ) except RuntimeError as e: logger.exception("Error in reasoning parser creation.") @@ -1052,10 +1080,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # wasn't ready to send a token, then # get the next token without streaming a chunk if delta_message is None: - if output.finish_reason is None: + if output.finish_reason is None and not request.return_token_ids: continue - else: - delta_message = DeltaMessage() + delta_message = DeltaMessage() # Log streaming delta if output logging is enabled if self.enable_log_outputs and self.request_logger: @@ -1438,13 +1465,20 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if self.reasoning_parser: try: - reasoning_parser = self.reasoning_parser(tokenizer) + chat_template_kwargs = self._prepare_extra_chat_template_kwargs( + request.chat_template_kwargs, + self.default_chat_template_kwargs, + ) + reasoning_parser = self.reasoning_parser( + tokenizer, + chat_template_kwargs=chat_template_kwargs, # type: ignore + ) except RuntimeError as e: logger.exception("Error in reasoning parser creation.") return self.create_error_response(str(e)) # If the reasoning parser is enabled, # tool calls are extracted exclusively from the content. - reasoning_content, content = reasoning_parser.extract_reasoning_content(output.text, request=request) + reasoning_content, content = reasoning_parser.extract_reasoning(output.text, request=request) if not request.include_reasoning: reasoning_content = None else: @@ -2049,7 +2083,9 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) -> ErrorResponse: """Create an error response following OpenAI error format.""" return ErrorResponse( - message=message, - type=err_type, - code=status_code, + error=ErrorInfo( + message=message, + type=err_type, + code=status_code, + ) ) 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 37733e290..45a1447b3 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 @@ -107,6 +107,7 @@ class Qwen2_5OmniForConditionalGeneration( ) if t2w_token_end_id: self.model.set_suppress_start_id(t2w_token_end_id + 1) + self.requires_raw_input_tokens = True elif self.model_stage == "code2wav": self.thinker = None @@ -125,6 +126,7 @@ class Qwen2_5OmniForConditionalGeneration( self._token2wav_conds: dict[str, torch.Tensor] = {} self._token2wav_ref_mels: dict[str, torch.Tensor] = {} self.model = self.token2wav + self.requires_raw_input_tokens = True else: raise ValueError("Invalid model stage") @@ -248,14 +250,27 @@ class Qwen2_5OmniForConditionalGeneration( if is_npu(): # TODO: remove this hack when NPU supports batched inputs properly thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids - thinker_positions = positions[0] if positions.ndim > 1 else positions + # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it + if positions.ndim == 2 and positions.shape[0] == 3: + thinker_positions = positions # MRoPE positions, keep as is + else: + thinker_positions = positions[0] if positions.ndim > 1 else positions thinker_inputs_embeds = ( inputs_embeds[0] if inputs_embeds is not None and added_batch_dim else inputs_embeds ) else: - thinker_input_ids = input_ids - thinker_positions = positions[0] - thinker_inputs_embeds = inputs_embeds + # Squeeze back if we added batch dim earlier + thinker_input_ids = input_ids[0] if input_ids is not None and added_batch_dim else input_ids + # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it + if positions.ndim == 2 and positions.shape[0] == 3: + thinker_positions = positions # MRoPE positions, keep as is + elif added_batch_dim: + thinker_positions = positions[0] + else: + thinker_positions = positions + thinker_inputs_embeds = ( + inputs_embeds[0] if inputs_embeds is not None and added_batch_dim else inputs_embeds + ) # Run thinker thinker_output = self.thinker( @@ -288,10 +303,16 @@ class Qwen2_5OmniForConditionalGeneration( if not hasattr(self, "voice_type"): self.voice_type = voice_type + # For MRoPE, positions shape is [3, num_tokens] (T/H/W), don't slice it + if positions.ndim == 2 and positions.shape[0] == 3: + talker_positions = positions # MRoPE positions, keep as is + else: + talker_positions = positions[0] + with torch.inference_mode(): talker_hidden = self.talker( input_ids=input_ids, - positions=positions[0], + positions=talker_positions, inputs_embeds=inputs_embeds, ) 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 71c5e8377..a2250b82e 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,4 +1,4 @@ -"""Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker (v0.12) with minimal overrides.""" +"""Thin Omni wrapper: reuse upstream Qwen2.5-Omni thinker (v0.14) with minimal overrides.""" from collections.abc import Iterable from typing import Any @@ -12,6 +12,7 @@ from transformers.models.qwen2_5_omni.modeling_qwen2_5_omni import ( Qwen2_5OmniAudioEncoder, ) from vllm.config import VllmConfig +from vllm.forward_context import set_forward_context from vllm.logger import init_logger from vllm.model_executor.models.interfaces import ( MultiModalEmbeddings, @@ -26,8 +27,6 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import ( 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, @@ -46,7 +45,9 @@ from vllm.model_executor.models.utils import ( WeightsMapper, init_vllm_registered_model, maybe_prefix, + split_list_into_ranges, ) +from vllm.model_executor.models.vision import get_llm_pos_ids_for_vision from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( MultiModalFeatureSpec, @@ -166,6 +167,43 @@ class Qwen2_5OmniConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixi video_grid_thw=video_grid_thw, ) + 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) + with set_forward_context(None, self.vllm_config): + 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) + with set_forward_context(None, self.vllm_config): + 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, @@ -180,8 +218,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration( SupportsMRoPE, Qwen2_5OmniConditionalGenerationMixin, ): - merge_by_field_config = True - hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "thinker.lm_head.": "language_model.lm_head.", @@ -250,6 +286,7 @@ class Qwen2_5OmniThinkerForConditionalGeneration( norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), quant_config=quant_config, prefix=maybe_prefix(prefix, "visual"), + multimodal_config=multimodal_config, ) else: self.visual = None 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 ce9819b3e..e04010196 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,9 +3,6 @@ from collections.abc import Iterable import torch from torch import nn from transformers import Qwen2Config -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 @@ -15,7 +12,6 @@ from vllm.model_executor.layers.activation import SiluAndMul from vllm.model_executor.layers.layernorm import RMSNorm from vllm.model_executor.layers.linear import MergedColumnParallelLinear, QKVParallelLinear, RowParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor -from vllm.model_executor.layers.pooler import Pooler, PoolingType from vllm.model_executor.layers.quantization import QuantizationConfig from vllm.model_executor.layers.rotary_embedding import get_rope from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead, VocabParallelEmbedding @@ -24,15 +20,14 @@ from vllm.model_executor.models.interfaces import SupportsLoRA, SupportsPP from vllm.model_executor.models.utils import ( AutoWeightsLoader, PPMissingLayer, - WeightsMapper, is_pp_missing_parameter, make_empty_intermediate_tensors_factory, make_layers, maybe_prefix, ) from vllm.sequence import IntermediateTensors -from vllm.v1.outputs import PoolerOutput, SamplerOutput -from vllm.v1.pool.metadata import PoolingMetadata +from vllm.v1.attention.backend import AttentionType +from vllm.v1.outputs import SamplerOutput from vllm.v1.sample.metadata import SamplingMetadata from vllm.v1.sample.sampler import Sampler @@ -130,7 +125,6 @@ class Qwen2Attention(nn.Module): self.rotary_pos_emb = get_rope( head_size=self.head_dim, - rotary_dim=self.head_dim, max_position=max_position, is_neox_style=True, rope_parameters={ @@ -460,66 +454,3 @@ class Qwen2ForCausalLM(nn.Module, SupportsLoRA, SupportsPP): skip_prefixes=(["lm_head."] if self.config.tie_word_embeddings else None), ) return loader.load_weights(weights) - - -class Qwen2EmbeddingModel(nn.Module, SupportsLoRA, SupportsPP): - packed_modules_mapping = { - "qkv_proj": [ - "q_proj", - "k_proj", - "v_proj", - ], - "gate_up_proj": [ - "gate_proj", - "up_proj", - ], - } - - hf_to_vllm_mapper = WeightsMapper(orig_to_new_prefix={"model.": ""}) - - def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""): - super().__init__() - config = vllm_config.model_config.hf_config - quant_config = vllm_config.quant_config - lora_config = vllm_config.lora_config - pooler_config = vllm_config.model_config.pooler_config - - self.config = config - self.lora_config = lora_config - - self.quant_config = quant_config - self.model = Qwen2Model(vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")) - - # TODO: Replace this model class with as_embedding_model( - # Qwen2ForCausalLM) after changing the default pooling method - if pooler_config.pooling_type is None: - logger.warning( - "This embedding model will default to last-token pooling in " - "an upcoming version. To avoid breaking changes, you should " - 'pass `--override-pooler-config \'{"pooling_type": "MEAN"}\'`' - " explicitly." - ) - - self._pooler = Pooler.from_config_with_defaults( - pooler_config, pooling_type=PoolingType.MEAN, normalize=True, softmax=False - ) - - def forward( - self, - input_ids: torch.Tensor, - positions: torch.Tensor, - intermediate_tensors: IntermediateTensors | None = None, - ) -> torch.Tensor: - return self.model(input_ids, positions, intermediate_tensors) - - def pooler( - self, - hidden_states: torch.Tensor, - pooling_metadata: PoolingMetadata, - ) -> PoolerOutput | None: - return self._pooler(hidden_states, pooling_metadata) - - def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]): - weights = self.hf_to_vllm_mapper.apply(weights) - weights = ((name, data) for name, data in weights if not name.startswith("lm_head.")) - self.model.load_weights(weights) 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 ba46d4a14..17f16d1d9 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -154,6 +154,7 @@ class Qwen3OmniMoeForConditionalGeneration( # for CI: Initialize special tokens embeddings early to avoid AttributeError when loading dummy weights self._init_special_tokens_embeddings() + self.requires_raw_input_tokens = True elif self.model_stage == "code2wav": self.thinker = None @@ -168,6 +169,7 @@ class Qwen3OmniMoeForConditionalGeneration( architectures=["Qwen3OmniMoeCode2Wav"], ) self.model = self.code2wav + self.requires_raw_input_tokens = True else: raise ValueError( f"Invalid model_stage: {self.model_stage}. Must be one of: 'thinker', 'talker', 'code2wav'" @@ -354,18 +356,24 @@ class Qwen3OmniMoeForConditionalGeneration( elif self.model_stage == "code2wav": # Extract codec codes from input codes = [] - if input_ids is not None: + if input_ids.shape[0] % 16 == 0: codes.append(input_ids.reshape(1, 16, -1)) - else: - # for profile, we use max length from inputs_embeds - codes.append( - torch.zeros( - (1, 16, inputs_embeds.shape[1]), - dtype=torch.long, - device=inputs_embeds.device, - ) + logger.warning( + ( + "Input_ids length: %s is not divisible by 16, padding " + "with zeros. This should only happen in warm up." + ), + input_ids.shape[0], ) + input_ids_flatten = input_ids.reshape(-1) + input_ids_flatten = torch.cat( + [ + input_ids_flatten, + torch.zeros(16 - input_ids.shape[0] % 16, dtype=torch.long, device=input_ids.device), + ] + ) + codes.append(input_ids_flatten.reshape(1, 16, -1)) # Generate audio from codec codes audio_tensors = [] 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 361a9349b..2d479062e 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 @@ -37,9 +37,6 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeConfig, Qwen3OmniMoeThinkerConfig, ) -from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import ( - Qwen3OmniMoeAudioEncoder, -) from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import ( Qwen3OmniMoeProcessor, ) @@ -56,6 +53,7 @@ from vllm.model_executor.models.interfaces import ( 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, @@ -69,6 +67,7 @@ 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, + Qwen3OmniMoeAudioEncoder, _get_feat_extract_output_lengths, ) from vllm.model_executor.models.utils import ( @@ -248,11 +247,36 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( # https://github.com/huggingface/transformers/pull/41473 mm_kwargs = dict(mm_kwargs) tok_kwargs = dict(tok_kwargs) + mm_kwargs["audio_kwargs"] = dict(mm_kwargs.get("audio_kwargs") or {}) + mm_kwargs["text_kwargs"] = dict(mm_kwargs.get("text_kwargs") or {}) if Version(TRANSFORMERS_VERSION) < Version("4.58.0"): + # Extract audio_sample_rate before restructuring + audio_sample_rate = mm_kwargs.pop("audio_sample_rate", None) + # move truncation to audio_kwargs level to avoid conflict # with tok_kwargs - mm_kwargs["audio_kwargs"] = {"truncation": mm_kwargs.pop("truncation", False)} - mm_kwargs["text_kwargs"] = {"truncation": tok_kwargs.pop("truncation", False)} + mm_kwargs["audio_kwargs"].setdefault("truncation", mm_kwargs.pop("truncation", False)) + mm_kwargs["text_kwargs"].setdefault("truncation", tok_kwargs.pop("truncation", False)) + + # Validate and conditionally pass audio_sample_rate + # WhisperFeatureExtractor has a fixed sampling rate, and vLLM's + # audio loader already resamples audio to the target rate. + # Only pass the value if it matches to avoid unexpected behavior. + if audio_sample_rate is not None: + expected_sr = feature_extractor.sampling_rate + if audio_sample_rate != expected_sr: + logger.warning( + "[%s] audio_sample_rate mismatch: user provided %dHz " + "but model expects %dHz. Ignoring user value. " + "vLLM's audio loader already resampled to %dHz.", + self.__class__.__name__, + audio_sample_rate, + expected_sr, + expected_sr, + ) + else: + # Sample rate matches, safe to pass + mm_kwargs["audio_kwargs"]["audio_sample_rate"] = audio_sample_rate hf_inputs = super()._call_hf_processor( prompt=prompt, @@ -398,11 +422,11 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( if audio_feature_lengths is None and feature_attention_mask is None: audio_output_lengths = [] elif audio_feature_lengths is not None: - _, audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths) + audio_output_lens = _get_feat_extract_output_lengths(audio_feature_lengths) audio_output_lengths = audio_output_lens.tolist() elif feature_attention_mask is not None: assert isinstance(feature_attention_mask, torch.Tensor) - _, audio_output_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1)) + audio_output_lens = _get_feat_extract_output_lengths(feature_attention_mask.sum(-1)) audio_output_lengths = audio_output_lens.tolist() # number of audios read from video. @@ -569,59 +593,21 @@ class Qwen3OmniMoeThinkerMultiModalProcessor( class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixin): - 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 - if ( - input_audio_features is not None - and isinstance(input_audio_features, torch.Tensor) - and input_audio_features.ndim == 3 - ): - # (batch_size, feature_dim, chunk_size) -> (feature_dim, batch_size * chunk_size) - input_audio_features = input_audio_features.permute(1, 0, 2).flatten(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): - for i in range(len(feature_attention_mask)): - feature_attention_mask[i] = feature_attention_mask[i].reshape(-1) - return Qwen2_5OmniAudioFeatureInputs( - type="audio_features", - input_features=input_audio_features, - audio_feature_lengths=audio_feature_lengths, - feature_attention_mask=feature_attention_mask, - ) - def _process_audio_input( self, audio_input: Qwen2_5OmniAudioFeatureInputs, audio_hashes: list[str] | None = None, cached_audio_features: torch.Tensor | None = None, - ) -> torch.Tensor: + ) -> tuple[torch.Tensor, ...]: input_features = audio_input["input_features"] audio_feature_lengths = audio_input["audio_feature_lengths"] - audio_feat_lengths, audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths) + audio_output_lengths = _get_feat_extract_output_lengths(audio_feature_lengths) audio_outputs = self.audio_tower( input_features.to(self.audio_tower.dtype), feature_lens=audio_feature_lengths, + aftercnn_lens=audio_output_lengths, ) audio_features = audio_outputs.last_hidden_state return audio_features.split(audio_output_lengths.tolist()) @@ -639,8 +625,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( SupportsMRoPE, Qwen3OmniMoeConditionalGenerationMixin, ): - merge_by_field_config = True - hf_to_vllm_mapper = WeightsMapper( orig_to_new_prefix={ "thinker.lm_head.": "language_model.lm_head.", @@ -649,6 +633,18 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( } ) + packed_modules_mapping = { + "qkv_proj": [ + "q_proj", + "k_proj", + "v_proj", + ], + "gate_up_proj": [ + "gate_proj", + "up_proj", + ], + } + @classmethod def get_placeholder_str(cls, modality: str, i: int) -> str | None: if modality.startswith("image"): @@ -669,28 +665,16 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( self.config = thinker_config self.multimodal_config = multimodal_config - # force "use_flash_attention_2=True" to audio tower to align - # the results. - if flash_attn is not None: - audio_config = thinker_config.audio_config - audio_config._attn_implementation_autoset = True - audio_config._attn_implementation = "flash_attention_2" - else: - logger.warning( - "flash_attn is not available, the model may not yield the " - "exactly same result as the transformers implementation " - "in the audio tower part." - ) + self.audio_tower = Qwen3OmniMoeAudioEncoder( + thinker_config.audio_config, + ) - self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config) - - attn_backend_override = multimodal_config.mm_encoder_attn_backend if multimodal_config is not None else None self.visual = Qwen3Omni_VisionTransformer( vision_config=thinker_config.vision_config, norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6), quant_config=quant_config, prefix=maybe_prefix(prefix, "visual"), - attn_backend_override=attn_backend_override, + multimodal_config=multimodal_config, ) self.quant_config = quant_config @@ -811,8 +795,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( return inputs_embeds deepstack_input_embeds = None - # TODO (ywang96): support overlapping modalitiy embeddings so that - # `use_audio_in_video` will work on V1. # split the feat dim to obtain multi-scale visual feature has_vision_embeddings = [ embeddings.shape[-1] != self.config.text_config.hidden_size for embeddings in multimodal_embeddings @@ -1006,7 +988,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( bos_len = 1 llm_pos_ids_list.append(torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx) st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - _, audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx]) + audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx]) llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx llm_pos_ids_list.append(llm_pos_ids) st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 @@ -1078,7 +1060,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( llm_pos_ids_list.append(bos_block) llm_pos_ids_list.append(bos_block) st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0 - _, audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx]) + audio_len = _get_feat_extract_output_lengths(audio_feature_lengths[audio_idx]) audio_llm_pos_ids = torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1) + st_idx grid_t = video_grid_thw[video_idx][0] grid_hs = video_grid_thw[:, 1] @@ -1122,3 +1104,13 @@ class Qwen3OmniMoeThinkerForConditionalGeneration( mrope_position_delta = llm_positions.max() + 1 - seq_len return llm_positions, mrope_position_delta + + def get_mm_mapping(self) -> MultiModelKeys: + """ + Get the module prefix in multimodal models + """ + return MultiModelKeys.from_string_field( + language_model="language_model", + connector="visual.merger", + tower_model=["visual.", "audio_tower."], + ) diff --git a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml index fc84f485e..e3d740ad5 100644 --- a/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen2_5_omni.yaml @@ -76,6 +76,7 @@ stage_args: trust_remote_code: true enable_prefix_caching: false max_num_batched_tokens: 32768 + async_scheduling: false engine_output_type: audio engine_input_source: [1] final_output: true diff --git a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml index c63dc5638..d4de07823 100644 --- a/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml +++ b/vllm_omni/model_executor/stage_configs/qwen3_omni_moe.yaml @@ -16,7 +16,7 @@ stage_args: worker_cls: vllm_omni.worker.gpu_ar_worker.GPUARWorker scheduler_cls: vllm_omni.core.sched.omni_ar_scheduler.OmniARScheduler gpu_memory_utilization: 0.6 - enforce_eager: false + enforce_eager: true trust_remote_code: true engine_output_type: latent # Output hidden states for talker distributed_executor_backend: "mp" @@ -46,8 +46,8 @@ 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.3 - enforce_eager: false + gpu_memory_utilization: 0.35 + enforce_eager: true trust_remote_code: true engine_output_type: latent # Output codec codes for code2wav # tensor_parallel_size: 2 @@ -80,6 +80,7 @@ stage_args: scheduler_cls: vllm_omni.core.sched.omni_generation_scheduler.OmniGenerationScheduler enforce_eager: true trust_remote_code: true + async_scheduling: false enable_prefix_caching: false engine_output_type: audio # Final output: audio waveform gpu_memory_utilization: 0.1 diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 2d2b7ef8e..18c8bc761 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -12,10 +12,15 @@ from typing import Any, NamedTuple import numpy as np import torch from vllm.config import CUDAGraphMode +from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer +from vllm.distributed.kv_transfer import get_kv_transfer_group from vllm.forward_context import set_forward_context from vllm.logger import init_logger +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, +) from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput -from vllm.v1.outputs import AsyncModelRunnerOutput +from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output 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 @@ -27,6 +32,7 @@ from vllm.v1.worker.gpu_model_runner import ( get_tp_group, has_kv_transfer_group, ) +from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import is_residual_scattered_for_sp from vllm_omni.outputs import OmniModelRunnerOutput @@ -44,6 +50,7 @@ class ExecuteModelState(NamedTuple): sample_hidden_states: torch.Tensor aux_hidden_states: list[torch.Tensor] | None ec_connector_output: Any + cudagraph_stats: Any multimodal_outputs: Any @@ -82,67 +89,138 @@ class GPUARModelRunner(OmniGPUModelRunner): scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors | None: - with record_function_or_nullcontext("Preprocess"): - with self.synchronize_input_prep(): - self._update_states(scheduler_output) - self._decode_and_store_request_payloads(scheduler_output) + if self.execute_model_state is not None: + raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") - if not scheduler_output.total_num_scheduled_tokens: - if not has_kv_transfer_group(): - return EMPTY_MODEL_RUNNER_OUTPUT - return self.kv_connector_no_forward(scheduler_output, self.vllm_config) - if self.cache_config.kv_sharing_fast_prefill: - assert not self.input_batch.num_prompt_logprobs, ( - "--kv-sharing-fast-prefill produces incorrect " - "logprobs for prompt tokens, tokens, please disable " - "it when the requests need prompt logprobs" - ) + if self.vllm_config.model_config.enable_return_routed_experts: + capturer = RoutedExpertsCapturer.get_instance() + if capturer is not None: + capturer.clear_buffer() # noqa + else: + logger.error("RoutedExpertsCapturer not initialized.") - 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 + if scheduler_output.preempted_req_ids and has_kv_transfer_group(): + get_kv_transfer_group().handle_preemptions(scheduler_output.preempted_req_ids) - logits_indices, spec_decode_metadata = self._prepare_inputs( + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + with ( + record_function_or_nullcontext("gpu_model_runner: preprocess"), + self.synchronize_input_prep(), + ): + # Update persistent batch states. + self._update_states(scheduler_output) + + if has_ec_transfer() and get_ec_transfer().is_producer: + with self.maybe_get_ec_connector_output( scheduler_output, + encoder_cache=self.encoder_cache, + ) as ec_connector_output: + self._execute_mm_encoder(scheduler_output) + return make_empty_encoder_model_runner_output(scheduler_output) + + if not num_scheduled_tokens: + if ( + self.parallel_config.distributed_executor_backend == "external_launcher" + and self.parallel_config.data_parallel_size > 1 + ): + # this is a corner case when both external launcher + # and DP are enabled, num_scheduled_tokens could be + # 0, and has_unfinished_requests in the outer loop + # returns True. before returning early here we call + # dummy run to ensure coordinate_batch_across_dp + # is called into to avoid out of sync issues. + self._dummy_run(1) + 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: + assert not self.num_prompt_logprobs, ( + "--kv-sharing-fast-prefill produces incorrect " + "logprobs for prompt tokens, tokens, please disable " + "it when the requests need prompt logprobs" + ) + + 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, + ) + + cascade_attn_prefix_lens = None + # Disable cascade attention when using microbatching (DBO) + if self.cascade_attn_enabled and not self.parallel_config.use_ubatching: + # Pre-compute cascade attention prefix lengths + cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens( num_scheduled_tokens_np, + self.input_batch.num_computed_tokens_cpu[:num_reqs], + scheduler_output.num_common_prefix_blocks, ) - ( - 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, - ) + ( + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + cudagraph_stats, + ) = 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=cascade_attn_prefix_lens is not None, + num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), + ) - 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 + logger.debug( + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " + "should_ubatch: %s, num_tokens_across_dp: %s", + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + ) - ( - attn_metadata, - spec_decode_common_attn_metadata, - ) = 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_tokens_padded = batch_desc.num_tokens + num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( + should_ubatch, + num_scheduled_tokens_np, + num_tokens_padded, + num_reqs_padded, + self.parallel_config.num_ubatches, + ) + + logger.debug( + "ubatch_slices: %s, ubatch_slices_padded: %s", + ubatch_slices, + ubatch_slices_padded, + ) + + pad_attn = cudagraph_mode == CUDAGraphMode.FULL + + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + + attn_metadata, spec_decode_common_attn_metadata = 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_attn, + logits_indices=logits_indices, + use_spec_decode=use_spec_decode, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + cascade_attn_prefix_lens=cascade_attn_prefix_lens, + ) ( input_ids, @@ -151,16 +229,18 @@ class GPUARModelRunner(OmniGPUModelRunner): intermediate_tensors, model_kwargs, ec_connector_output, - ) = self._preprocess( - scheduler_output, - num_tokens_padded, - intermediate_tensors, - ) + ) = self._preprocess(scheduler_output, num_tokens_padded, intermediate_tensors) + # Set cudagraph mode to none if calc_kv_scales is true. + # KV scales calculation involves dynamic operations that are incompatible + # with CUDA graph capture. if self.calculate_kv_scales: cudagraph_mode = CUDAGraphMode.NONE + # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False + # Run the model. + # Use persistent buffers for CUDA graphs. with ( set_forward_context( attn_metadata, @@ -169,9 +249,9 @@ class GPUARModelRunner(OmniGPUModelRunner): num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=cudagraph_mode, batch_descriptor=batch_desc, - ubatch_slices=ubatch_slices, + ubatch_slices=ubatch_slices_padded, ), - record_function_or_nullcontext("Forward"), + record_function_or_nullcontext("gpu_model_runner: forward"), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, ): model_output = self._model_forward( @@ -208,27 +288,32 @@ class GPUARModelRunner(OmniGPUModelRunner): 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 + self.kv_connector_output = kv_connector_output return hidden_states if self.is_pooling_model: - output = self._pool( + # Return the pooling output. + return self._pool( hidden_states, - num_tokens_padded, + num_scheduled_tokens, num_scheduled_tokens_np, + kv_connector_output, ) - output.kv_connector_output = kv_connector_output - return output sample_hidden_states = hidden_states[logits_indices] logits = self.model.compute_logits( sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata ) else: + # Rare case. assert not self.is_pooling_model + sample_hidden_states = hidden_states[logits_indices] if not get_pp_group().is_last_rank: all_gather_tensors = { "residual": not is_residual_scattered_for_sp(self.vllm_config, num_tokens_padded) @@ -240,7 +325,6 @@ class GPUARModelRunner(OmniGPUModelRunner): ) logits = None else: - sample_hidden_states = hidden_states[logits_indices] logits = self.model.compute_logits( sample_hidden_states, sampling_metadata=self.input_batch.sampling_metadata ) @@ -264,6 +348,7 @@ class GPUARModelRunner(OmniGPUModelRunner): sample_hidden_states, aux_hidden_states, ec_connector_output, + cudagraph_stats, multimodal_outputs, ) self.kv_connector_output = kv_connector_output @@ -278,14 +363,20 @@ class GPUARModelRunner(OmniGPUModelRunner): self.kv_connector_output = None if self.execute_model_state is None: + # Nothing to do (PP non-final rank case), output isn't used. if not kv_connector_output: return None # type: ignore[return-value] + + # In case of PP with kv transfer, we need to pass through the + # kv_connector_output 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 + # Unpack ephemeral state. ( scheduler_output, logits, @@ -295,16 +386,20 @@ class GPUARModelRunner(OmniGPUModelRunner): sample_hidden_states, aux_hidden_states, ec_connector_output, + cudagraph_stats, multimodal_outputs, ) = self.execute_model_state self.execute_model_state = None + # Apply structured output bitmasks if present. 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._draft_token_ids = None + self._draft_token_req_ids = None self.input_batch.prev_sampled_token_ids = None def propose_draft_token_ids(sampled_token_ids): @@ -320,39 +415,40 @@ class GPUARModelRunner(OmniGPUModelRunner): spec_decode_metadata, spec_decode_common_attn_metadata, ) + self._copy_draft_token_ids_to_cpu(scheduler_output) spec_config = self.speculative_config - use_padded_batch_for_eagle = ( - 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 ( - 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 = spec_config.draft_model_config.max_model_len - input_fits_in_drafter = spec_decode_common_attn_metadata and ( - spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens <= effective_drafter_max_model_len - ) - 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) + propose_drafts_after_bookkeeping = False + if spec_config is not None: + input_fits_in_drafter = spec_decode_common_attn_metadata is not None and ( + spec_decode_common_attn_metadata.max_seq_len + self.num_spec_tokens + <= self.effective_drafter_max_model_len + ) + if spec_config.use_eagle() and not spec_config.disable_padded_drafter_batch: + # EAGLE speculative decoding can use the GPU sampled tokens + # as inputs, and does not need to wait for bookkeeping to finish. + 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) + # Since we couldn't run the drafter, + # just use zeros for the draft tokens. + self._draft_token_ids = torch.zeros(1, device=self.device, dtype=torch.int32).expand( + len(self.input_batch.req_ids), self.num_spec_tokens + ) + self._copy_draft_token_ids_to_cpu(scheduler_output, zeros_only=True) + else: + propose_drafts_after_bookkeeping = input_fits_in_drafter with record_function_or_nullcontext("gpu_model_runner: bookkeep"): ( @@ -372,7 +468,9 @@ class GPUARModelRunner(OmniGPUModelRunner): spec_decode_metadata, ) - if self.speculative_config and not use_padded_batch_for_eagle and input_fits_in_drafter: + if propose_drafts_after_bookkeeping: + # 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("gpu_model_runner: eplb"): @@ -421,6 +519,12 @@ class GPUARModelRunner(OmniGPUModelRunner): payload.update(mm_payload) pooler_output.append(payload) with record_function_or_nullcontext("gpu_model_runner: ModelRunnerOutput"): + if self.model_config.enable_return_routed_experts: + capturer = RoutedExpertsCapturer.get_instance() + if capturer is not None: + capturer.save_captured_experts(indices=self.slot_mapping) # noqa + else: + logger.error("RoutedExpertsCapturer not initialized.") output = OmniModelRunnerOutput( req_ids=req_ids_output_copy, req_id_to_index=req_id_to_index_output_copy, @@ -431,6 +535,7 @@ class GPUARModelRunner(OmniGPUModelRunner): 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, + cudagraph_stats=cudagraph_stats, ) if not self.use_async_scheduling: diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 9e058addb..d2dafce68 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -3,13 +3,14 @@ import os import torch from vllm.logger import init_logger -from vllm.model_executor import set_random_seed from vllm.platforms import current_platform -from vllm.utils.mem_constants import GiB_bytes -from vllm.utils.mem_utils import MemorySnapshot +from vllm.utils.mem_utils import MemorySnapshot, format_gib +from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment +from vllm.v1.worker.utils import request_memory +from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner @@ -24,21 +25,19 @@ class GPUARWorker(GPUWorker): """ def init_device(self): - device = self.device_config.device - if isinstance(device, torch.device) and device.type == "cuda": + if self.device_config.device_type == "cuda": # This env var set by Ray causes exceptions with graph building. os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + parallel_config = self.parallel_config 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 + parallel_config.distributed_executor_backend not in ("ray", "external_launcher") + and parallel_config.data_parallel_backend != "ray" + and 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 + dp_local_rank = self.parallel_config.data_parallel_index tp_pp_world_size = ( self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size @@ -80,25 +79,17 @@ class GPUARWorker(GPUWorker): torch.cuda.empty_cache() # take current memory snapshot - self.init_snapshot = MemorySnapshot() - self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization - if self.init_snapshot.free_memory < self.requested_memory: - - def gib(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"is less than desired GPU memory utilization " - f"({self.cache_config.gpu_memory_utilization}, " - f"{gib(self.requested_memory)} GiB). Decrease GPU memory " - f"utilization or reduce GPU memory used by other processes." - ) + self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) + self.requested_memory = request_memory(init_snapshot, self.cache_config) + logger.debug("worker init memory snapshot: %r", self.init_snapshot) + logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory)) else: raise RuntimeError(f"Not support device type: {self.device_config.device}") + # Initialize workspace manager + num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 + init_workspace_manager(self.device, num_ubatches) + # Construct the model runner self.model_runner = GPUARModelRunner(self.vllm_config, self.device) diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 17740d858..011c74d2e 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -8,13 +8,20 @@ from __future__ import annotations import gc import logging +from copy import copy import numpy as np import torch from vllm.config import CUDAGraphMode -from vllm.multimodal.inputs import MultiModalKwargs +from vllm.distributed.ec_transfer import get_ec_transfer, has_ec_transfer +from vllm.distributed.kv_transfer import get_kv_transfer_group, has_kv_transfer_group +from vllm.model_executor.layers.fused_moe.routed_experts_capturer import ( + RoutedExpertsCapturer, +) +from vllm.model_executor.models.interfaces import supports_mm_encoder_only from vllm.utils.math_utils import cdiv -from vllm.v1.core.sched.output import SchedulerOutput +from vllm.v1.core.sched.output import GrammarOutput, SchedulerOutput +from vllm.v1.outputs import AsyncModelRunnerOutput, make_empty_encoder_model_runner_output from vllm.v1.spec_decode.eagle import EagleProposer from vllm.v1.utils import record_function_or_nullcontext from vllm.v1.worker.gpu_model_runner import ( @@ -25,9 +32,11 @@ from vllm.v1.worker.gpu_model_runner import ( get_pp_group, set_forward_context, ) +from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs from vllm_omni.outputs import OmniModelRunnerOutput +from vllm_omni.worker.gpu_ar_model_runner import ExecuteModelState from vllm_omni.worker.gpu_model_runner import OmniGPUModelRunner logger = logging.getLogger(__name__) @@ -47,57 +56,139 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): scheduler_output: SchedulerOutput, intermediate_tensors: IntermediateTensors | None = None, ) -> OmniModelRunnerOutput | IntermediateTensors: - with record_function_or_nullcontext("Preprocess"): - with self.synchronize_input_prep(): - self._update_states(scheduler_output) - if not scheduler_output.total_num_scheduled_tokens: + if self.execute_model_state is not None: + raise RuntimeError("State error: sample_tokens() must be called after execute_model() returns None.") + + if self.vllm_config.model_config.enable_return_routed_experts: + capturer = RoutedExpertsCapturer.get_instance() + if capturer is not None: + capturer.clear_buffer() # noqa + else: + logger.error("RoutedExpertsCapturer not initialized.") + + if scheduler_output.preempted_req_ids and has_kv_transfer_group(): + get_kv_transfer_group().handle_preemptions(scheduler_output.preempted_req_ids) + + num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens + with ( + record_function_or_nullcontext("gpu_model_runner: preprocess"), + self.synchronize_input_prep(), + ): + self._update_states(scheduler_output) + if not scheduler_output.total_num_scheduled_tokens: + return EMPTY_MODEL_RUNNER_OUTPUT + + if has_ec_transfer() and get_ec_transfer().is_producer: + with self.maybe_get_ec_connector_output( + scheduler_output, + encoder_cache=self.encoder_cache, + ) as ec_connector_output: + self._execute_mm_encoder(scheduler_output) + return make_empty_encoder_model_runner_output(scheduler_output) + + if not num_scheduled_tokens: + if ( + self.parallel_config.distributed_executor_backend == "external_launcher" + and self.parallel_config.data_parallel_size > 1 + ): + # this is a corner case when both external launcher + # and DP are enabled, num_scheduled_tokens could be + # 0, and has_unfinished_requests in the outer loop + # returns True. before returning early here we call + # dummy run to ensure coordinate_batch_across_dp + # is called into to avoid out of sync issues. + self._dummy_run(1) + if not has_kv_transfer_group(): + # Return empty ModelRunnerOutput if no work to do. return EMPTY_MODEL_RUNNER_OUTPUT - 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 + return self.kv_connector_no_forward(scheduler_output, self.vllm_config) - logits_indices, spec_decode_metadata = self._prepare_inputs( - scheduler_output, + if self.cache_config.kv_sharing_fast_prefill: + assert not self.num_prompt_logprobs, ( + "--kv-sharing-fast-prefill produces incorrect " + "logprobs for prompt tokens, tokens, please disable " + "it when the requests need prompt logprobs" + ) + 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, + ) + + cascade_attn_prefix_lens = None + # Disable cascade attention when using microbatching (DBO) + if self.cascade_attn_enabled and not self.parallel_config.use_ubatching: + # Pre-compute cascade attention prefix lengths + cascade_attn_prefix_lens = self._compute_cascade_attn_prefix_lens( num_scheduled_tokens_np, + self.input_batch.num_computed_tokens_cpu[:num_reqs], + scheduler_output.num_common_prefix_blocks, ) - ( - 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, - ) + ( + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + cudagraph_stats, + ) = 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=cascade_attn_prefix_lens is not None, + num_encoder_reqs=len(scheduler_output.scheduled_encoder_inputs), + ) - 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 + logger.debug( + "Running batch with cudagraph_mode: %s, batch_descriptor: %s, " + "should_ubatch: %s, num_tokens_across_dp: %s", + cudagraph_mode, + batch_desc, + should_ubatch, + num_tokens_across_dp, + ) - ( - attn_metadata, - spec_decode_common_attn_metadata, - ) = 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_tokens_padded = batch_desc.num_tokens + num_reqs_padded = batch_desc.num_reqs if batch_desc.num_reqs is not None else num_reqs + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( + should_ubatch, + num_scheduled_tokens_np, + num_tokens_padded, + num_reqs_padded, + self.parallel_config.num_ubatches, + ) + + logger.debug( + "ubatch_slices: %s, ubatch_slices_padded: %s", + ubatch_slices, + ubatch_slices_padded, + ) + + pad_attn = cudagraph_mode == CUDAGraphMode.FULL + + use_spec_decode = len(scheduler_output.scheduled_spec_decode_tokens) > 0 + ubatch_slices_attn = ubatch_slices_padded if pad_attn else ubatch_slices + + attn_metadata, spec_decode_common_attn_metadata = 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_attn, + logits_indices=logits_indices, + use_spec_decode=use_spec_decode, + num_scheduled_tokens=scheduler_output.num_scheduled_tokens, + cascade_attn_prefix_lens=cascade_attn_prefix_lens, + ) ( input_ids, @@ -112,10 +203,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): intermediate_tensors, ) + # Set cudagraph mode to none if calc_kv_scales is true. + # KV scales calculation involves dynamic operations that are incompatible + # with CUDA graph capture. if self.calculate_kv_scales: cudagraph_mode = CUDAGraphMode.NONE + # Mark KV scales as calculated after the first forward pass self.calculate_kv_scales = False + # Run the model. + # Use persistent buffers for CUDA graphs. with ( set_forward_context( attn_metadata, @@ -124,7 +221,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=cudagraph_mode, batch_descriptor=batch_desc, - ubatch_slices=ubatch_slices, + ubatch_slices=ubatch_slices_padded, ), record_function_or_nullcontext("Forward"), self.maybe_get_kv_connector_output(scheduler_output) as kv_connector_output, @@ -139,6 +236,61 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): ) _, multimodal_outputs = self.extract_multimodal_outputs(outputs) + self.execute_model_state = ExecuteModelState( + scheduler_output, + None, + spec_decode_metadata, + spec_decode_common_attn_metadata, + None, + None, + None, + ec_connector_output, + cudagraph_stats, + multimodal_outputs, + ) + self.kv_connector_output = kv_connector_output + return None + + @torch.inference_mode() + def sample_tokens( + self, + grammar_output: GrammarOutput | None = None, + ) -> OmniModelRunnerOutput | AsyncModelRunnerOutput | IntermediateTensors: + # NOTE: Even though the model is non-autoregressive, we still need + # this function to match the interface of the engine core. + # In this case, this function + kv_connector_output = self.kv_connector_output + self.kv_connector_output = None + + if self.execute_model_state is None: + # Nothing to do (PP non-final rank case), output isn't used. + if not kv_connector_output: + return None # type: ignore[return-value] + + # In case of PP with kv transfer, we need to pass through the + # kv_connector_output + 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 + + # Unpack ephemeral state. + ( + scheduler_output, + logits, + spec_decode_metadata, + spec_decode_common_attn_metadata, + hidden_states, + sample_hidden_states, + aux_hidden_states, + ec_connector_output, + cudagraph_stats, + multimodal_outputs, + ) = self.execute_model_state + self.execute_model_state = None + pooler_output: list[object] = [] if isinstance(multimodal_outputs, torch.Tensor): assert multimodal_outputs.shape[0] == 1, ( @@ -169,6 +321,8 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): pooler_output=pooler_output, kv_connector_output=kv_connector_output, num_nans_in_logits={}, + cudagraph_stats=cudagraph_stats, + ec_connector_output=ec_connector_output if self.supports_mm_inputs else None, ) if not self.use_async_scheduling: @@ -176,9 +330,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): return AsyncGPUModelRunnerOutput( model_runner_output=output, - sampled_token_ids=[], + sampled_token_ids=torch.tensor([], device=self.device), invalid_req_indices=[], async_output_copy_stream=self.async_output_copy_stream, + vocab_size=self.input_batch.vocab_size, + logprobs_tensors=None, ) def _run_generation_model( @@ -206,7 +362,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): positions=positions, intermediate_tensors=intermediate_tensors, inputs_embeds=inputs_embeds, - **MultiModalKwargs.as_kwargs(model_kwargs, device=self.device), + **model_kwargs, sampling_metadata=self.input_batch.sampling_metadata, logits_index=logits_indices, sampler=self.sampler, @@ -263,6 +419,11 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): remove_lora: If False, dummy LoRAs are not destroyed after the run activate_lora: If False, dummy_run is performed without LoRAs. """ + if supports_mm_encoder_only(self.model): + # The current dummy run only covers LM execution, so we can skip it. + # mm encoder dummy run may need to add in the future. + return torch.tensor([]), torch.tensor([]) + assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes() # If cudagraph_mode.decode_mode() == FULL and @@ -316,23 +477,25 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) - _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, + _cudagraph_mode, batch_desc, should_ubatch, 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: @@ -345,6 +508,18 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): 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 + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( + should_ubatch, + num_scheduled_tokens, + num_tokens_padded, + num_reqs_padded, + self.vllm_config.parallel_config.num_ubatches, + ) + logger.debug( + "ubatch_slices: %s, ubatch_slices_padded: %s", + ubatch_slices, + ubatch_slices_padded, + ) attn_metadata: PerLayerAttnMetadata | None = None @@ -366,11 +541,12 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens self.query_start_loc.copy_to_gpu() + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL 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, + ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, ) @@ -383,10 +559,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): ): # 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) + model_kwargs = self._init_model_kwargs() if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: - input_ids = None - inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded) + model_kwargs = { **model_kwargs, **self._dummy_mm_kwargs(num_reqs), @@ -394,7 +570,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): elif self.enable_prompt_embeds: input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] - model_kwargs = self._init_model_kwargs(num_tokens_padded) + model_kwargs = self._init_model_kwargs() else: input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None @@ -418,16 +594,16 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False) - if ubatch_slices is not None: + if ubatch_slices_padded 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 + num_tokens_padded = ubatch_slices_padded[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), + self.maybe_randomize_inputs(input_ids, inputs_embeds), set_forward_context( attn_metadata, self.vllm_config, @@ -435,7 +611,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, - ubatch_slices=ubatch_slices, + ubatch_slices=ubatch_slices_padded, ), ): outputs = self.model( @@ -453,10 +629,13 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): 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) + # Eagle currently only supports PIECEWISE cudagraphs. + # Therefore only use cudagraphs if the main model uses PIECEWISE + # NOTE(lucas): this is a hack, need to clean up. use_cudagraphs = ( - cudagraph_runtime_mode.has_mode(CUDAGraphMode.PIECEWISE) - and not self.speculative_config.enforce_eager - ) + (is_graph_capturing and cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE) + or (not is_graph_capturing and cudagraph_runtime_mode != CUDAGraphMode.NONE) + ) 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 @@ -471,6 +650,17 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): is_graph_capturing=is_graph_capturing, ) + # We register layerwise NVTX hooks here after the first dynamo tracing is + # done to avoid nvtx operations in hook functions being traced by + # torch dynamo and causing graph breaks. + # Note that for DYNAMO_ONCE and VLLM_COMPILE mode, + # compiled model's dynamo tracing is only done once and the compiled model's + # __call__ function is replaced by calling the compiled function. + # So it's safe to register hooks here. Hooks will be registered to + # both compiled and uncompiled models but they will never + # be called on the compiled model execution path. + self._register_layerwise_nvtx_hooks() + # This is necessary to avoid blocking DP. # For dummy runs, we typically skip EPLB since we don't have any real # requests to process. diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index 27111f394..19f8ab84b 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -2,16 +2,20 @@ import gc import os import torch -from vllm.model_executor import set_random_seed +from vllm.logger import init_logger from vllm.platforms import current_platform -from vllm.utils.mem_constants import GiB_bytes -from vllm.utils.mem_utils import MemorySnapshot +from vllm.utils.mem_utils import MemorySnapshot, format_gib +from vllm.utils.torch_utils import set_random_seed from vllm.v1.utils import report_usage_stats from vllm.v1.worker.gpu_worker import Worker as GPUWorker from vllm.v1.worker.gpu_worker import init_worker_distributed_environment +from vllm.v1.worker.utils import request_memory +from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner +logger = init_logger(__name__) + class GPUGenerationWorker(GPUWorker): """GPU Worker for Generation model (non-autoregressive waveform generation). @@ -21,21 +25,19 @@ class GPUGenerationWorker(GPUWorker): """ def init_device(self): - device = self.device_config.device - if isinstance(device, torch.device) and device.type == "cuda": + if self.device_config.device_type == "cuda": # This env var set by Ray causes exceptions with graph building. os.environ.pop("NCCL_ASYNC_ERROR_HANDLING", None) + parallel_config = self.parallel_config 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 + parallel_config.distributed_executor_backend not in ("ray", "external_launcher") + and parallel_config.data_parallel_backend != "ray" + and 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 + dp_local_rank = self.parallel_config.data_parallel_index tp_pp_world_size = ( self.parallel_config.pipeline_parallel_size * self.parallel_config.tensor_parallel_size @@ -77,26 +79,17 @@ class GPUGenerationWorker(GPUWorker): torch.cuda.empty_cache() # take current memory snapshot - self.init_snapshot = MemorySnapshot() - self.requested_memory = self.init_snapshot.total_memory * self.cache_config.gpu_memory_utilization - if self.init_snapshot.free_memory < self.requested_memory: - - def gib(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"is less than desired GPU memory utilization " - f"({self.cache_config.gpu_memory_utilization}, " - f"{gib(self.requested_memory)} GiB). Decrease GPU memory " - f"utilization or reduce GPU memory used by other processes." - ) + self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) + self.requested_memory = request_memory(init_snapshot, self.cache_config) + logger.debug("worker init memory snapshot: %r", self.init_snapshot) + logger.debug("worker requested memory: %sGiB", format_gib(self.requested_memory)) else: raise RuntimeError(f"Not support device type: {self.device_config.device}") - # Construct the model runner + # Initialize workspace manager + num_ubatches = 2 if self.vllm_config.parallel_config.enable_dbo else 1 + init_workspace_manager(self.device, num_ubatches) + self.model_runner = GPUGenerationModelRunner(self.vllm_config, self.device) if self.rank == 0: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index e10301a52..06abd9f3a 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -8,7 +8,7 @@ from vllm.distributed.parallel_state import get_pp_group 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 import supports_mm_encoder_only, supports_mrope from vllm.model_executor.models.interfaces_base import VllmModelForPooling from vllm.sampling_params import SamplingType from vllm.utils.import_utils import LazyLoader @@ -16,6 +16,7 @@ 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_utils import maybe_create_ubatch_slices from vllm_omni.model_executor.models.output_templates import OmniOutput @@ -129,6 +130,7 @@ class OmniGPUModelRunner(GPUModelRunner): # Remove finished requests from the cached states. for req_id in scheduler_output.finished_req_ids: self.requests.pop(req_id, None) + self.num_prompt_logprobs.pop(req_id, None) # Remove the finished requests from the persistent batch. # NOTE(woosuk): There could be an edge case where finished_req_ids and # scheduled_req_ids overlap. This happens when a request is aborted and @@ -149,7 +151,14 @@ class OmniGPUModelRunner(GPUModelRunner): # they will be scheduled again sometime in the future. scheduled_req_ids = scheduler_output.num_scheduled_tokens.keys() cached_req_ids = self.input_batch.req_id_to_index.keys() - unscheduled_req_ids = cached_req_ids - scheduled_req_ids + resumed_req_ids = scheduler_output.scheduled_cached_reqs.resumed_req_ids + # NOTE(zhuohan): cached_req_ids and resumed_req_ids are usually disjoint, + # so `(scheduled_req_ids - resumed_req_ids) == scheduled_req_ids` holds + # apart from the forced-preemption case in reset_prefix_cache. And in + # that case we include the resumed_req_ids in the unscheduled set so + # that they get cleared from the persistent batch before being re-scheduled + # in the normal resumed request path. + unscheduled_req_ids = cached_req_ids - (scheduled_req_ids - resumed_req_ids) # NOTE(woosuk): The persistent batch optimization assumes that # consecutive batches contain mostly the same requests. If batches # have low request overlap (e.g., alternating between two distinct @@ -241,20 +250,62 @@ class OmniGPUModelRunner(GPUModelRunner): logger.error(f"Error decoding additional information: {e}") pass + if sampling_params and sampling_params.prompt_logprobs is not None: + self.num_prompt_logprobs[req_id] = ( + self.input_batch.vocab_size + if sampling_params.prompt_logprobs == -1 + else sampling_params.prompt_logprobs + ) # Only relevant for models using M-RoPE (e.g, Qwen2-VL) if self.uses_mrope: self._init_mrope_positions(req_state) + # Only relevant for models using XD-RoPE (e.g, HunYuan-VL) + if self.uses_xdrope_dim > 0: + self._init_xdrope_positions(req_state) + reqs_to_add.append(self.requests[req_id]) # Update the states of the running/resumed requests. is_last_rank = get_pp_group().is_last_rank req_data = scheduler_output.scheduled_cached_reqs + scheduled_spec_tokens = scheduler_output.scheduled_spec_decode_tokens + + # Wait until valid_sampled_tokens_count is copied to cpu, + # then use it to update actual num_computed_tokens of each request. + valid_sampled_token_count = self._get_valid_sampled_token_count() + for i, req_id in enumerate(req_data.req_ids): req_state = self.requests[req_id] num_computed_tokens = req_data.num_computed_tokens[i] new_block_ids = req_data.new_block_ids[i] - resumed_from_preemption = req_data.resumed_from_preemption[i] + resumed_from_preemption = req_id in req_data.resumed_req_ids + num_output_tokens = req_data.num_output_tokens[i] + req_index = self.input_batch.req_id_to_index.get(req_id) + + if req_state.prev_num_draft_len and self.use_async_scheduling: + # prev_num_draft_len is used in async scheduling mode with + # spec decode. it indicates if need to update num_computed_tokens + # of the request. for example: + # fist step: num_computed_tokens = 0, spec_tokens = [], + # prev_num_draft_len = 0. + # second step: num_computed_tokens = 100(prompt length), + # spec_tokens = [a,b], prev_num_draft_len = 0. + # third step: num_computed_tokens = 100 + 2, spec_tokens = [c,d], + # prev_num_draft_len = 2. + # num_computed_tokens in first step and second step does't contain + # the spec tokens length, but in third step it contains the + # spec tokens length. we only need to update num_computed_tokens + # when prev_num_draft_len > 0. + if req_index is None: + req_state.prev_num_draft_len = 0 + else: + assert self.input_batch.prev_req_id_to_index is not None + prev_req_index = self.input_batch.prev_req_id_to_index[req_id] + num_accepted = valid_sampled_token_count[prev_req_index] - 1 + num_rejected = req_state.prev_num_draft_len - num_accepted + num_computed_tokens -= num_rejected + req_state.output_token_ids.extend([-1] * num_accepted) # Update the cached states. req_state.num_computed_tokens = num_computed_tokens @@ -272,6 +323,13 @@ class OmniGPUModelRunner(GPUModelRunner): req_state.output_token_ids.append(new_token_ids[-1]) elif num_new_tokens > 0: req_state.output_token_ids.extend(new_token_ids[-num_new_tokens:]) + elif num_output_tokens < len(req_state.output_token_ids): + # Some output tokens were discarded due to a sync-KV-load + # failure. Align the cached state. + del req_state.output_token_ids[num_output_tokens:] + if req_index is not None: + end_idx = self.input_batch.num_prompt_tokens[req_index] + num_output_tokens + self.input_batch.num_tokens_no_spec[req_index] = end_idx # Update the block IDs. if not resumed_from_preemption: @@ -280,6 +338,7 @@ class OmniGPUModelRunner(GPUModelRunner): for block_ids, new_ids in zip(req_state.block_ids, new_block_ids): block_ids.extend(new_ids) else: + assert req_index is None assert new_block_ids is not None # The request is resumed from preemption. # Replace the existing block IDs with the new ones. @@ -290,6 +349,13 @@ class OmniGPUModelRunner(GPUModelRunner): # The request is not in the persistent batch. # The request was either preempted and resumed later, or was not # scheduled in the previous step and needs to be added again. + + if self.use_async_scheduling and num_output_tokens > 0: + # We must recover the output token ids for resumed requests in the + # async scheduling case, so that correct input_ids are obtained. + resumed_token_ids = req_data.all_token_ids[req_id] + req_state.output_token_ids = resumed_token_ids[-num_output_tokens:] + reqs_to_add.append(req_state) continue @@ -306,22 +372,15 @@ class OmniGPUModelRunner(GPUModelRunner): end_token_index = num_computed_tokens + len(new_token_ids) self.input_batch.token_ids_cpu[req_index, start_token_index:end_token_index] = new_token_ids self.input_batch.num_tokens_no_spec[req_index] = end_token_index - self.input_batch.num_tokens[req_index] = end_token_index # Add spec_token_ids to token_ids_cpu. - spec_token_ids = scheduler_output.scheduled_spec_decode_tokens.get(req_id, ()) - if spec_token_ids: - num_spec_tokens = len(spec_token_ids) - start_index = self.input_batch.num_tokens_no_spec[req_index] - end_token_index = start_index + num_spec_tokens - self.input_batch.token_ids_cpu[req_index, start_index:end_token_index] = spec_token_ids - # NOTE(woosuk): `num_tokens` here may include spec tokens. - self.input_batch.num_tokens[req_index] += num_spec_tokens + self.input_batch.update_req_spec_token_ids(req_state, scheduled_spec_tokens) # Add the new or resumed requests to the persistent batch. # The smaller empty indices are filled first. for request in reqs_to_add: self.input_batch.add_request(request) + self.input_batch.update_req_spec_token_ids(request, scheduled_spec_tokens) # Condense the batched states if there are gaps left by removed requests self.input_batch.condense() @@ -388,6 +447,11 @@ class OmniGPUModelRunner(GPUModelRunner): remove_lora: If False, dummy LoRAs are not destroyed after the run activate_lora: If False, dummy_run is performed without LoRAs. """ + if supports_mm_encoder_only(self.model): + # The current dummy run only covers LM execution, so we can skip it. + # mm encoder dummy run may need to add in the future. + return torch.tensor([]), torch.tensor([]) + assert cudagraph_runtime_mode is None or cudagraph_runtime_mode.valid_runtime_modes() # If cudagraph_mode.decode_mode() == FULL and @@ -441,23 +505,25 @@ class OmniGPUModelRunner(GPUModelRunner): num_sampled_tokens = np.ones(num_reqs, dtype=np.int32) - _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, + _cudagraph_mode, batch_desc, should_ubatch, 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: @@ -470,6 +536,18 @@ class OmniGPUModelRunner(GPUModelRunner): 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 + ubatch_slices, ubatch_slices_padded = maybe_create_ubatch_slices( + should_ubatch, + num_scheduled_tokens, + num_tokens_padded, + num_reqs_padded, + self.vllm_config.parallel_config.num_ubatches, + ) + logger.debug( + "ubatch_slices: %s, ubatch_slices_padded: %s", + ubatch_slices, + ubatch_slices_padded, + ) attn_metadata: PerLayerAttnMetadata | None = None @@ -491,11 +569,12 @@ class OmniGPUModelRunner(GPUModelRunner): self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens self.query_start_loc.copy_to_gpu() + pad_attn = cudagraph_runtime_mode == CUDAGraphMode.FULL 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, + ubatch_slices=ubatch_slices_padded if pad_attn else ubatch_slices, for_cudagraph_capture=is_graph_capturing, ) @@ -508,10 +587,10 @@ class OmniGPUModelRunner(GPUModelRunner): ): # 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) + model_kwargs = self._init_model_kwargs() if self.supports_mm_inputs and not self.model_config.is_encoder_decoder: - input_ids = None - inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] + input_ids, inputs_embeds = self._prepare_mm_inputs(num_tokens_padded) + model_kwargs = { **model_kwargs, **self._dummy_mm_kwargs(num_reqs), @@ -519,7 +598,7 @@ class OmniGPUModelRunner(GPUModelRunner): elif self.enable_prompt_embeds: input_ids = None inputs_embeds = self.inputs_embeds.gpu[:num_tokens_padded] - model_kwargs = self._init_model_kwargs(num_tokens_padded) + model_kwargs = self._init_model_kwargs() else: input_ids = self.input_ids.gpu[:num_tokens_padded] inputs_embeds = None @@ -543,16 +622,16 @@ class OmniGPUModelRunner(GPUModelRunner): intermediate_tensors = self.sync_and_slice_intermediate_tensors(num_tokens_padded, None, False) - if ubatch_slices is not None: + if ubatch_slices_padded 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 + num_tokens_padded = ubatch_slices_padded[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), + self.maybe_randomize_inputs(input_ids, inputs_embeds), set_forward_context( attn_metadata, self.vllm_config, @@ -560,7 +639,7 @@ class OmniGPUModelRunner(GPUModelRunner): num_tokens_across_dp=num_tokens_across_dp, cudagraph_runtime_mode=cudagraph_runtime_mode, batch_descriptor=batch_desc, - ubatch_slices=ubatch_slices, + ubatch_slices=ubatch_slices_padded, ), ): if getattr(self.model, "talker", None) is not None and hasattr(self.model, "talker_mtp"): @@ -585,10 +664,13 @@ class OmniGPUModelRunner(GPUModelRunner): 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) + # Eagle currently only supports PIECEWISE cudagraphs. + # Therefore only use cudagraphs if the main model uses PIECEWISE + # NOTE(lucas): this is a hack, need to clean up. use_cudagraphs = ( - cudagraph_runtime_mode.has_mode(CUDAGraphMode.PIECEWISE) - and not self.speculative_config.enforce_eager - ) + (is_graph_capturing and cudagraph_runtime_mode == CUDAGraphMode.PIECEWISE) + or (not is_graph_capturing and cudagraph_runtime_mode != CUDAGraphMode.NONE) + ) 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 @@ -603,6 +685,17 @@ class OmniGPUModelRunner(GPUModelRunner): is_graph_capturing=is_graph_capturing, ) + # We register layerwise NVTX hooks here after the first dynamo tracing is + # done to avoid nvtx operations in hook functions being traced by + # torch dynamo and causing graph breaks. + # Note that for DYNAMO_ONCE and VLLM_COMPILE mode, + # compiled model's dynamo tracing is only done once and the compiled model's + # __call__ function is replaced by calling the compiled function. + # So it's safe to register hooks here. Hooks will be registered to + # both compiled and uncompiled models but they will never + # be called on the compiled model execution path. + self._register_layerwise_nvtx_hooks() + # This is necessary to avoid blocking DP. # For dummy runs, we typically skip EPLB since we don't have any real # requests to process. @@ -772,7 +865,7 @@ class OmniGPUModelRunner(GPUModelRunner): num_input_tokens: int, intermediate_tensors: IntermediateTensors | None = None, ): - """Align with v0.12 preprocess and omni's additional information handling.""" + """Align with v0.14.0 preprocess and omni's additional information handling.""" 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 @@ -802,10 +895,9 @@ class OmniGPUModelRunner(GPUModelRunner): # 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] + input_ids, inputs_embeds = self._prepare_mm_inputs(num_input_tokens) model_kwargs = { - **self._init_model_kwargs(num_scheduled_tokens), + **self._init_model_kwargs(), **self._extract_mm_kwargs(scheduler_output), } elif self.enable_prompt_embeds and is_first_rank: @@ -829,7 +921,7 @@ class OmniGPUModelRunner(GPUModelRunner): 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) + model_kwargs = self._init_model_kwargs() input_ids = self.input_ids.gpu[:num_input_tokens] else: # For text-only models, we use token ids as input. @@ -838,7 +930,7 @@ class OmniGPUModelRunner(GPUModelRunner): # 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) + model_kwargs = self._init_model_kwargs() if self.uses_mrope: positions = self.mrope_positions.gpu[:, :num_input_tokens] @@ -929,7 +1021,7 @@ class OmniGPUModelRunner(GPUModelRunner): decode_batch_size = len(decode_req_ids) if decode_batch_size == 0: return - _cudagraph_mode, batch_desc, _, _ = self._determine_batch_execution_and_padding( + _cudagraph_mode, batch_desc, _, _, _ = self._determine_batch_execution_and_padding( num_tokens=decode_batch_size, num_reqs=decode_batch_size, num_scheduled_tokens_np=np.ones(decode_batch_size, dtype=np.int32),