From 2804a85e58875f559267d2f811ca2b9e78dd2dae Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Sat, 4 Apr 2026 15:03:32 +0800 Subject: [PATCH] [Rebase] Rebase to vllm v0.19.0 (#2475) --- .buildkite/pipeline-intel.yaml | 2 +- .buildkite/test-nightly.yml | 2 +- docker/Dockerfile.ci | 6 +- docker/Dockerfile.rocm | 2 +- docker/Dockerfile.xpu | 2 +- .../installation/gpu/cuda.inc.md | 11 +-- .../installation/gpu/rocm.inc.md | 14 ++-- docs/getting_started/quickstart.md | 4 +- docs/mkdocs/hooks/generate_argparse.py | 1 + .../quantization/test_quantization_quality.py | 8 +- .../offline_inference/test_bagel_text2img.py | 20 +++++ tests/e2e/online_serving/test_mimo_audio.py | 31 +++++--- tests/engine/test_async_omni_engine_abort.py | 1 + .../openai_api/test_serving_speech.py | 1 + tests/entrypoints/test_async_omni_abort.py | 4 +- tests/entrypoints/test_omni_entrypoints.py | 1 + tests/examples/conftest.py | 3 +- .../models/test_omni_processing.py | 9 ++- vllm_omni/benchmarks/patch/patch.py | 8 +- vllm_omni/core/sched/omni_ar_scheduler.py | 5 +- .../core/sched/omni_generation_scheduler.py | 7 +- .../diffusion/attention/backends/utils/fa.py | 4 +- vllm_omni/diffusion/diffusion_engine.py | 4 +- vllm_omni/diffusion/stage_diffusion_client.py | 15 +++- vllm_omni/diffusion/stage_diffusion_proc.py | 2 +- .../diffusion/worker/diffusion_worker.py | 10 ++- vllm_omni/engine/arg_utils.py | 48 +++++++++++- vllm_omni/engine/async_omni_engine.py | 40 +++++++++- vllm_omni/engine/stage_init_utils.py | 29 +++++++ vllm_omni/entrypoints/async_omni.py | 13 +++- vllm_omni/entrypoints/cli/main.py | 4 + vllm_omni/entrypoints/cli/serve.py | 32 ++++++++ vllm_omni/entrypoints/openai/api_server.py | 7 +- vllm_omni/entrypoints/openai/serving_chat.py | 36 ++++----- .../entrypoints/openai/serving_speech.py | 3 + vllm_omni/entrypoints/utils.py | 23 ++++-- vllm_omni/inputs/data.py | 6 +- vllm_omni/inputs/preprocess.py | 17 ++-- .../model_executor/models/bagel/bagel.py | 2 +- .../models/cosyvoice3/cosyvoice3.py | 3 +- .../models/glm_image/glm_image_ar.py | 2 +- .../models/hunyuan_image3/hunyuan_image3.py | 2 +- .../models/mimo_audio/mimo_audio.py | 3 +- .../models/mimo_audio/mimo_audio_llm.py | 2 +- .../models/qwen3_omni/qwen3_omni.py | 2 +- .../qwen3_omni/qwen3_omni_moe_thinker.py | 2 +- .../voxtral_tts_audio_generation.py | 2 +- vllm_omni/patch.py | 2 +- .../npu/worker/npu_ar_model_runner.py | 4 +- vllm_omni/worker/gpu_ar_model_runner.py | 16 +++- .../worker/gpu_generation_model_runner.py | 25 ++++-- vllm_omni/worker/gpu_model_runner.py | 77 ++++++++++++++----- 52 files changed, 437 insertions(+), 142 deletions(-) diff --git a/.buildkite/pipeline-intel.yaml b/.buildkite/pipeline-intel.yaml index 4334dd516..2dc53ad96 100644 --- a/.buildkite/pipeline-intel.yaml +++ b/.buildkite/pipeline-intel.yaml @@ -10,7 +10,7 @@ steps: DOCKER_BUILDKIT: "1" # Buildkite will automatically replace this with the actual commit hash VLLM_IMAGE_TAG: "${BUILDKITE_COMMIT}" - VLLM_VERSION: "v0.18.0" + VLLM_VERSION: "v0.19.0" priority: 100 timeout_in_minutes: 60 soft_fail: true diff --git a/.buildkite/test-nightly.yml b/.buildkite/test-nightly.yml index 32bf219bc..9dc888506 100644 --- a/.buildkite/test-nightly.yml +++ b/.buildkite/test-nightly.yml @@ -362,7 +362,7 @@ steps: - kubernetes: podSpec: containers: - - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:c392ce21e9cf9ea65c52b866447793db10e0261c + - image: 936637512419.dkr.ecr.us-west-2.amazonaws.com/vllm-ci-pull-through-cache/q9t5s3a7/vllm-ci-test-repo:$BUILDKITE_COMMIT resources: limits: nvidia.com/gpu: 2 diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index f4253fe72..24ce39baf 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.18.0 +ARG VLLM_BASE_TAG=v0.19.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} @@ -11,9 +11,7 @@ RUN apt-get update && \ apt-get clean && \ rm -rf /var/lib/apt/lists/* -# Install vllm-omni into the same uv-managed Python environment used by the base image. -# Use bash -c so that $(python3 -c ...) is expanded inside the container. -RUN uv pip install --system --no-cache-dir ".[dev]" +RUN uv pip install --system ".[dev]" RUN ln -sf /usr/bin/python3 /usr/bin/python diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index ce541497a..bfbb060bc 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.18.0 +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.19.0 FROM ${BASE_IMAGE} AS base # Declare a variable to know if we want to use the nightly build or the stable build. diff --git a/docker/Dockerfile.xpu b/docker/Dockerfile.xpu index 8901725f0..17f1aebf0 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -76,7 +76,7 @@ ENV UV_EXTRA_INDEX_URL=${PIP_EXTRA_INDEX_URL} ENV UV_INDEX_STRATEGY="unsafe-best-match" ENV UV_LINK_MODE="copy" -ARG VLLM_VERSION=v0.18.0 +ARG VLLM_VERSION=v0.19.0 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm WORKDIR /workspace/vllm diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index 932b54f0d..3e6600f66 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -39,13 +39,13 @@ uv pip install 'vllm-omni[demo]' # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library ```bash -uv pip install vllm==0.18.0 --torch-backend=auto +uv pip install vllm==0.19.0 --torch-backend=auto ``` -The 0.18.0 release of vLLM ships CUDA 12.9-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. +The 0.19.0 release of vLLM ships CUDA 13.0-compatible binaries by default. If you need a different CUDA variant or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -66,11 +66,12 @@ 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.18.0 +git checkout v0.19.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.18.0/vllm-0.18.0+cu129-cp38-abi3-manylinux_2_35_x86_64.whl +#For CUDA 13.0 +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.19.0/vllm-0.19.0+cu130-cp38-abi3-manylinux_2_35_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index da84561c9..1a683d174 100644 --- a/docs/getting_started/installation/gpu/rocm.inc.md +++ b/docs/getting_started/installation/gpu/rocm.inc.md @@ -13,7 +13,7 @@ vLLM-Omni current recommends the steps in under setup through Docker Images. vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700 +uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700 ``` #### Installation of vLLM-Omni @@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm sox # --8<-- [start:build-wheel-from-source] #### Installation of vLLM -If you do not need to modify source code of vLLM, you can directly install the stable 0.18.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.19.0 release version of the library ```bash -uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700 +uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700 ``` -The pre-built 0.18.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. +The pre-built 0.19.0 vLLM wheel targets ROCm 7.0. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. #### Installation of vLLM-Omni Since vllm-omni is rapidly evolving, it's recommended to install it from source @@ -58,7 +58,7 @@ 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.18.0 +git checkout v0.19.0 python3 -m pip install -r requirements/rocm.txt python3 setup.py develop ``` @@ -130,7 +130,7 @@ docker run --rm \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ - vllm/vllm-omni-rocm:v0.18.0 \ + vllm/vllm-omni-rocm:v0.19.0 \ --model Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` @@ -149,7 +149,7 @@ docker run --rm -it \ -v ~/.cache/huggingface:/root/.cache/huggingface \ --env "HF_TOKEN=$HF_TOKEN" \ --entrypoint bash \ - vllm/vllm-omni-rocm:v0.18.0 + vllm/vllm-omni-rocm:v0.19.0 ``` # --8<-- [end:pre-built-images] diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index eef3dd1a7..45b3eab1d 100644 --- a/docs/getting_started/quickstart.md +++ b/docs/getting_started/quickstart.md @@ -19,10 +19,10 @@ uv venv --python 3.12 --seed source .venv/bin/activate # On CUDA -uv pip install vllm==0.18.0 --torch-backend=auto +uv pip install vllm==0.19.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.18.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.18.0/rocm700 +uv pip install vllm==0.19.0+rocm700 --extra-index-url https://wheels.vllm.ai/rocm/0.19.0/rocm700 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni diff --git a/docs/mkdocs/hooks/generate_argparse.py b/docs/mkdocs/hooks/generate_argparse.py index 6cef7cfbd..0aed44a0c 100644 --- a/docs/mkdocs/hooks/generate_argparse.py +++ b/docs/mkdocs/hooks/generate_argparse.py @@ -120,6 +120,7 @@ def extract_omni_serve_subparser_init(): "_FlexibleArgumentParser": _FlexibleArgumentParser, "FlexibleArgumentParser": _FlexibleArgumentParser, "make_arg_parser": lambda parser: parser, # no-op for doc + "_ensure_vllm_platform": lambda: None, # no-op for doc "VLLM_SUBCMD_PARSER_EPILOG": "", "logger": logger, "DummySubparsers": DummySubparsers, diff --git a/tests/diffusion/quantization/test_quantization_quality.py b/tests/diffusion/quantization/test_quantization_quality.py index a937a6486..3d8f18736 100644 --- a/tests/diffusion/quantization/test_quantization_quality.py +++ b/tests/diffusion/quantization/test_quantization_quality.py @@ -118,8 +118,12 @@ def _generate_image(omni, config: QualityTestConfig): peak_mem = torch.cuda.max_memory_allocated() / (1024**3) first = outputs[0] - req_out = first.request_output[0] if hasattr(first, "request_output") else first - return req_out.images[0], peak_mem + if hasattr(first, "images") and first.images: + return first.images[0], peak_mem + inner = first.request_output + if inner is not None and hasattr(inner, "images") and inner.images: + return inner.images[0], peak_mem + raise ValueError("Could not extract image from output.") def _generate_video(omni, config: QualityTestConfig): diff --git a/tests/e2e/offline_inference/test_bagel_text2img.py b/tests/e2e/offline_inference/test_bagel_text2img.py index c74763a35..7cce8da3a 100644 --- a/tests/e2e/offline_inference/test_bagel_text2img.py +++ b/tests/e2e/offline_inference/test_bagel_text2img.py @@ -229,6 +229,24 @@ def _wait_for_port(host: str, port: int, timeout: int = 30) -> bool: return False +def _is_mooncake_master_available() -> bool: + """Check if mooncake_master binary is present and can actually execute.""" + import shutil + + binary = shutil.which("mooncake_master") + if binary is None: + return False + try: + result = subprocess.run( + [binary, "--help"], + capture_output=True, + timeout=5, + ) + return result.returncode != 127 + except (subprocess.TimeoutExpired, OSError): + return True + + def _cleanup_mooncake_processes(timeout_secs: int = 5) -> None: """Clean up any existing mooncake_master processes. @@ -292,6 +310,8 @@ def _load_mooncake_config(host: str, rpc_port: int, http_port: int) -> str: @hardware_test(res={"cuda": "H100"}) def test_bagel_text2img_mooncake_connector(run_level): """Test Bagel text2img with Mooncake connector for inter-stage communication.""" + if not _is_mooncake_master_available(): + pytest.skip("mooncake_master is not available or cannot execute (missing shared libraries like libibverbs)") MOONCAKE_HOST = "127.0.0.1" MOONCAKE_RPC_PORT = _find_free_port() MOONCAKE_HTTP_PORT = _find_free_port() diff --git a/tests/e2e/online_serving/test_mimo_audio.py b/tests/e2e/online_serving/test_mimo_audio.py index 2fb63c1e4..43eeb7733 100644 --- a/tests/e2e/online_serving/test_mimo_audio.py +++ b/tests/e2e/online_serving/test_mimo_audio.py @@ -63,20 +63,27 @@ def download_tokenizer(): # CI stage config for H100 / MI325 -stage_configs = [get_chunk_config()] -tokenizer_path = download_tokenizer() -os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path +# Guard module-level setup so test collection doesn't fail in environments +# where the model cache is read-only or models aren't available. +try: + stage_configs = [get_chunk_config()] + tokenizer_path = download_tokenizer() + os.environ["MIMO_AUDIO_TOKENIZER_PATH"] = tokenizer_path -# Create parameter combinations for model and stage config -test_params = [ - OmniServerParams( - model=model, - stage_config_path=stage_config, - server_args=["--chat-template", CHAT_TEMPLATE_PATH], + test_params = [ + OmniServerParams( + model=model, + stage_config_path=stage_config, + server_args=["--chat-template", CHAT_TEMPLATE_PATH], + ) + for model in models + for stage_config in stage_configs + ] +except Exception as exc: + pytest.skip( + f"MiMo-Audio online serving tests skipped: module setup failed ({type(exc).__name__}: {exc})", + allow_module_level=True, ) - for model in models - for stage_config in stage_configs -] def get_prompt(prompt_type="text_only"): diff --git a/tests/engine/test_async_omni_engine_abort.py b/tests/engine/test_async_omni_engine_abort.py index a99c522c3..34fdf45ea 100644 --- a/tests/engine/test_async_omni_engine_abort.py +++ b/tests/engine/test_async_omni_engine_abort.py @@ -60,6 +60,7 @@ async def generate( @pytest.mark.core_model @pytest.mark.omni +@pytest.mark.real_hf_config @hardware_test(res={"cuda": "L4", "rocm": "MI325"}, num_cards=1) @pytest.mark.asyncio async def test_abort(): diff --git a/tests/entrypoints/openai_api/test_serving_speech.py b/tests/entrypoints/openai_api/test_serving_speech.py index 83e4188c1..b140b7a04 100644 --- a/tests/entrypoints/openai_api/test_serving_speech.py +++ b/tests/entrypoints/openai_api/test_serving_speech.py @@ -1664,6 +1664,7 @@ class TestFishSpeechServing: assert all(allowed_special is None for _, _, allowed_special in tokenizer.calls) def test_build_fish_clone_prompt_normalizes_text_fields(self, fish_speech_server): + fish_speech_server._fish_speech_tokenizer = _FakeFishTokenizer() fish_speech_server._estimate_fish_prompt_len = MagicMock(return_value=123) request = OpenAICreateSpeechRequest( diff --git a/tests/entrypoints/test_async_omni_abort.py b/tests/entrypoints/test_async_omni_abort.py index 71f3e99fe..b34652162 100644 --- a/tests/entrypoints/test_async_omni_abort.py +++ b/tests/entrypoints/test_async_omni_abort.py @@ -13,8 +13,8 @@ def test_generate_accepts_request_after_repeated_cancellations(): submitted_request_ids = [] aborted_request_batches = [] - async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id): - del prompt, sampling_params_list, final_stage_id + async def fake_add_request_async(*, request_id, prompt, sampling_params_list, final_stage_id, **kwargs): + del prompt, sampling_params_list, final_stage_id, kwargs submitted_request_ids.append(request_id) async def fake_abort_async(request_ids): diff --git a/tests/entrypoints/test_omni_entrypoints.py b/tests/entrypoints/test_omni_entrypoints.py index 0aeb6158e..3cffcd37d 100644 --- a/tests/entrypoints/test_omni_entrypoints.py +++ b/tests/entrypoints/test_omni_entrypoints.py @@ -113,6 +113,7 @@ class FakeAsyncOmniEngine: sampling_params_list: list[Any] | None = None, final_stage_id: int = 0, arrival_time: float | None = None, + **kwargs: Any, ) -> None: msg = { "request_id": request_id, diff --git a/tests/examples/conftest.py b/tests/examples/conftest.py index a66db9040..137d15f16 100644 --- a/tests/examples/conftest.py +++ b/tests/examples/conftest.py @@ -14,7 +14,6 @@ from collections.abc import Callable from pathlib import Path from typing import Any, NamedTuple, cast -import mistune import pytest import torch from safetensors.torch import save_file @@ -62,6 +61,8 @@ class ReadmeSnippet(NamedTuple): readme_path: Path, skipif: ReadmeSnippetExtractionSkipPredicate | None = None, ) -> list["ReadmeSnippet"]: + import mistune + markdown = mistune.create_markdown(renderer="ast") tokens = markdown(readme_path.read_text(encoding="utf-8")) tokens = cast(list[dict[str, Any]], tokens) # mistune's AST renderer always produces a list, not a str diff --git a/tests/model_executor/models/test_omni_processing.py b/tests/model_executor/models/test_omni_processing.py index fe9c63b82..70a9ca0e9 100644 --- a/tests/model_executor/models/test_omni_processing.py +++ b/tests/model_executor/models/test_omni_processing.py @@ -21,9 +21,10 @@ from vllm.config.multimodal import ( ImageDummyOptions, VideoDummyOptions, ) -from vllm.multimodal import MULTIMODAL_REGISTRY, MultiModalDataDict +from vllm.inputs import MultiModalDataDict, MultiModalInput +from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.cache import MultiModalProcessorOnlyCache -from vllm.multimodal.inputs import MultiModalInputs, batched_tensors_equal +from vllm.multimodal.inputs import batched_tensors_equal from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config @@ -302,8 +303,8 @@ def test_omni_processing_correctness( def _assert_inputs_equal( - a: MultiModalInputs, - b: MultiModalInputs, + a: MultiModalInput, + b: MultiModalInput, *, ignore_mm_keys: set[str] | None = None, msg: str = "", diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 539af11f8..d8145c40b 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -358,6 +358,7 @@ async def benchmark( lora_modules: Iterable[str] | None, extra_headers: dict | None, extra_body: dict | None, + lora_assignment: Literal["random", "round-robin"] = "random", ramp_up_strategy: Literal["linear", "exponential"] | None = None, ramp_up_start_rps: int | None = None, ramp_up_end_rps: int | None = None, @@ -454,8 +455,11 @@ async def benchmark( print("Starting main benchmark run...") if lora_modules: - # For each input request, choose a LoRA module at random. - lora_modules = iter([random.choice(lora_modules) for _ in range(len(input_requests))]) + lora_modules_list = list(lora_modules) + if lora_assignment == "round-robin": + lora_modules = iter([lora_modules_list[i % len(lora_modules_list)] for i in range(len(input_requests))]) + else: + lora_modules = iter([random.choice(lora_modules_list) for _ in range(len(input_requests))]) if profile: print("Starting profiler...") diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index d49664161..af178d14d 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -235,7 +235,10 @@ class OmniARScheduler(VLLMScheduler): # These blocks contain externally computed tokens that failed to # load. Identify affected requests and adjust their computed token # count to trigger recomputation of the invalid blocks. - failed_kv_load_req_ids = self._handle_invalid_blocks(kv_connector_output.invalid_block_ids) + failed_kv_load_req_ids = self._handle_invalid_blocks( + kv_connector_output.invalid_block_ids, + num_scheduled_tokens, + ) # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, # the below loop can be a performance bottleneck. We should do our best diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index dded8f7aa..1c4356d4f 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -272,7 +272,7 @@ class OmniGenerationScheduler(VLLMScheduler): # KVTransfer: package metadata if self.connector is not None: - meta = self.connector.build_connector_meta(scheduler_output) + meta = self._build_kv_connector_meta(self.connector, scheduler_output) scheduler_output.kv_connector_metadata = meta # EC Connector: package metadata if self.ec_connector is not None: @@ -368,7 +368,10 @@ class OmniGenerationScheduler(VLLMScheduler): failed_kv_load_req_ids = None if kv_connector_output and getattr(kv_connector_output, "invalid_block_ids", None): - failed_kv_load_req_ids = self._handle_invalid_blocks(kv_connector_output.invalid_block_ids) + failed_kv_load_req_ids = self._handle_invalid_blocks( + kv_connector_output.invalid_block_ids, + num_scheduled_tokens, + ) # NOTE(woosuk): As len(num_scheduled_tokens) can be up to 1K or more, # the below loop can be a performance bottleneck. We should do our best diff --git a/vllm_omni/diffusion/attention/backends/utils/fa.py b/vllm_omni/diffusion/attention/backends/utils/fa.py index 1fd47790f..77596a103 100644 --- a/vllm_omni/diffusion/attention/backends/utils/fa.py +++ b/vllm_omni/diffusion/attention/backends/utils/fa.py @@ -32,7 +32,9 @@ if current_omni_platform.is_rocm(): pass elif current_omni_platform.is_xpu(): try: - from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func # noqa: F401 + from vllm._xpu_ops import xpu_ops # noqa: F401 + + flash_attn_varlen_func = xpu_ops.flash_attn_varlen_func except (ImportError, ModuleNotFoundError): pass elif current_omni_platform.is_musa(): diff --git a/vllm_omni/diffusion/diffusion_engine.py b/vllm_omni/diffusion/diffusion_engine.py index 05008d7e9..784da6175 100644 --- a/vllm_omni/diffusion/diffusion_engine.py +++ b/vllm_omni/diffusion/diffusion_engine.py @@ -361,8 +361,8 @@ class DiffusionEngine: def _dummy_run(self): """A dummy run to warm up the model.""" num_inference_steps = 1 - height = 1024 - width = 1024 + height = 512 + width = 512 if supports_image_input(self.od_config.model_class_name): # Provide a dummy image input if the model supports it color_format = image_color_format(self.od_config.model_class_name) diff --git a/vllm_omni/diffusion/stage_diffusion_client.py b/vllm_omni/diffusion/stage_diffusion_client.py index db13f99aa..77db2b1b9 100644 --- a/vllm_omni/diffusion/stage_diffusion_client.py +++ b/vllm_omni/diffusion/stage_diffusion_client.py @@ -244,7 +244,20 @@ class StageDiffusionClient: return self._output_queue.get_nowait() except asyncio.QueueEmpty: if not self._shutting_down and self._proc is not None and not self._proc.is_alive(): - raise RuntimeError(f"StageDiffusionProc died unexpectedly (exit code {self._proc.exitcode})") + exitcode = self._proc.exitcode + # One final drain – the last ZMQ frame may have arrived + # between the first drain and the is_alive() check. + self._drain_responses() + try: + return self._output_queue.get_nowait() + except asyncio.QueueEmpty: + pass + if exitcode is not None and exitcode > 128: + sig = exitcode - 128 + logger.warning("StageDiffusionProc was killed by signal %d; treating as external shutdown.", sig) + self._shutting_down = True + return None + raise RuntimeError(f"StageDiffusionProc died unexpectedly (exit code {exitcode})") return None async def abort_requests_async(self, request_ids: list[str]) -> None: diff --git a/vllm_omni/diffusion/stage_diffusion_proc.py b/vllm_omni/diffusion/stage_diffusion_proc.py index 8677da037..0a5fd3590 100644 --- a/vllm_omni/diffusion/stage_diffusion_proc.py +++ b/vllm_omni/diffusion/stage_diffusion_proc.py @@ -495,7 +495,7 @@ class StageDiffusionProc: nonlocal shutdown_requested if not shutdown_requested: shutdown_requested = True - raise SystemExit() + raise SystemExit(128 + signum) signal.signal(signal.SIGTERM, signal_handler) signal.signal(signal.SIGINT, signal_handler) diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index 6e1cabba0..ea4b9d96f 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -17,7 +17,7 @@ from typing import Any import torch import zmq -from vllm.config import CompilationConfig, VllmConfig, set_current_vllm_config +from vllm.config import CompilationConfig, DeviceConfig, VllmConfig, set_current_vllm_config from vllm.distributed.device_communicators.shm_broadcast import MessageQueue from vllm.logger import init_logger from vllm.utils.import_utils import resolve_obj_by_qualname @@ -113,8 +113,12 @@ class DiffusionWorker: self.device = current_omni_platform.get_torch_device(rank) current_omni_platform.set_device(self.device) - # Create vllm_config for parallel configuration - vllm_config = VllmConfig(compilation_config=CompilationConfig()) + # Create vllm_config for parallel configuration. Pass explicit device_config + # so DeviceConfig does not rely on current_platform in worker subprocesses. + vllm_config = VllmConfig( + compilation_config=CompilationConfig(), + device_config=DeviceConfig(device=self.device), + ) vllm_config.parallel_config.tensor_parallel_size = self.od_config.parallel_config.tensor_parallel_size vllm_config.parallel_config.data_parallel_size = self.od_config.parallel_config.data_parallel_size vllm_config.parallel_config.enable_expert_parallel = self.od_config.parallel_config.enable_expert_parallel diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index 5bc51043a..b66378926 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -1,6 +1,8 @@ import argparse import dataclasses +import json import os +import tempfile from dataclasses import dataclass, field from typing import Any @@ -135,6 +137,30 @@ class OmniEngineArgs(EngineArgs): self._omni_models_registered = True return True + def _patch_empty_hf_config(self, model_type: str) -> None: + """For models with empty config.json (e.g. CosyVoice3), create a + patched config in a temp directory with model_type set so that + transformers AutoConfig.from_pretrained can resolve the config class. + Sets self.hf_config_path to point to the patched directory.""" + try: + from transformers import PretrainedConfig + + config_dict, _ = PretrainedConfig.get_config_dict(self.model) + if config_dict.get("model_type"): + return # config.json already has model_type, no patching needed + except Exception: + return # can't load config, let vLLM handle the error + + # Create a temp dir with a patched config.json + temp_dir = tempfile.mkdtemp(prefix="omni_hf_config_") + config_dict["model_type"] = model_type + config_dict.setdefault("architectures", [self.model_arch]) + with open(os.path.join(temp_dir, "config.json"), "w") as f: + json.dump(config_dict, f) + self.hf_config_path = temp_dir + self._temp_config_dir = temp_dir + logger.info("Patched empty HF config with model_type=%s at %s", model_type, temp_dir) + def create_model_config(self) -> OmniModelConfig: """Create an OmniModelConfig from these engine arguments. Returns: @@ -164,6 +190,18 @@ class OmniEngineArgs(EngineArgs): if model_type is not None: self.hf_overrides.setdefault("model_type", model_type) + # For models whose HF config.json is empty or lacks model_type + # (e.g. CosyVoice3), AutoConfig.from_pretrained fails because it + # cannot determine which config class to use from the empty dict. + # hf_overrides alone is not enough since transformers reads + # model_type from config_dict before applying overrides. + # Workaround: create a patched config.json in a temp directory + # and point hf_config_path to it so vLLM reads model_type from it. + if not self.hf_config_path: + model_type = _ARCH_TO_MODEL_TYPE.get(self.model_arch) + if model_type is not None: + self._patch_empty_hf_config(model_type) + # Auto-detect tokenizer for models that store it in a subdirectory # rather than the root (e.g. CosyVoice3 uses CosyVoice-BlankEN/). if not self.tokenizer and self.model: @@ -200,7 +238,15 @@ class OmniEngineArgs(EngineArgs): logger.warning("Failed to download tokenizer subfolder: %s", e) # Build the vLLM config first, then use it to create the Omni config. - model_config = super().create_model_config() + try: + model_config = super().create_model_config() + finally: + # Clean up temp config dir if we created one + if hasattr(self, "_temp_config_dir"): + import shutil + + shutil.rmtree(self._temp_config_dir, ignore_errors=True) + del self._temp_config_dir omni_config = OmniModelConfig.from_vllm_model_config( model_config=model_config, diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index d9960ecba..092b341e4 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -17,7 +17,7 @@ import threading import time import uuid import weakref -from collections.abc import Sequence +from collections.abc import Mapping, Sequence from dataclasses import asdict from typing import TYPE_CHECKING, Any @@ -652,6 +652,12 @@ class AsyncOmniEngine: sampling_params_list: Sequence[Any] | None = None, final_stage_id: int = 0, arrival_time: float | None = None, + lora_request: Any = None, + tokenization_kwargs: dict[str, Any] | None = None, + trace_headers: Mapping[str, str] | None = None, + priority: int = 0, + data_parallel_rank: int | None = None, + reasoning_ended: bool | None = None, *, resumable: bool = False, message_type: str = "add_request", @@ -686,12 +692,20 @@ class AsyncOmniEngine: params=params, supported_tasks=self.supported_tasks, arrival_time=arrival_time, + lora_request=lora_request, + tokenization_kwargs=tokenization_kwargs, + trace_headers=trace_headers, + priority=priority, + data_parallel_rank=data_parallel_rank, resumable=resumable, ) # TODO (Peiqi): add this for Qwen3-TTS only. Other models don't have # additional_information field in the prompt. request = _upgrade_to_omni_request(request, prompt) + if reasoning_ended is not None: + request.reasoning_ended = reasoning_ended + # Restore external_req_id to the original user-facing request_id. # InputProcessor.process_inputs() renames request_id to an internal # UUID (saving the original in external_req_id), but then overwrites @@ -973,6 +987,12 @@ class AsyncOmniEngine: sampling_params_list: Sequence[Any] | None = None, final_stage_id: int = 0, arrival_time: float | None = None, + lora_request: Any = None, + tokenization_kwargs: dict[str, Any] | None = None, + trace_headers: Mapping[str, str] | None = None, + priority: int = 0, + data_parallel_rank: int | None = None, + reasoning_ended: bool | None = None, *, resumable: bool = False, ) -> None: @@ -990,6 +1010,12 @@ class AsyncOmniEngine: sampling_params_list=sampling_params_list, final_stage_id=final_stage_id, arrival_time=arrival_time, + lora_request=lora_request, + tokenization_kwargs=tokenization_kwargs, + trace_headers=trace_headers, + priority=priority, + data_parallel_rank=data_parallel_rank, + reasoning_ended=reasoning_ended, resumable=resumable, ) if self.request_queue is None: @@ -1013,6 +1039,12 @@ class AsyncOmniEngine: sampling_params_list: Sequence[Any] | None = None, final_stage_id: int = 0, arrival_time: float | None = None, + lora_request: Any = None, + tokenization_kwargs: dict[str, Any] | None = None, + trace_headers: Mapping[str, str] | None = None, + priority: int = 0, + data_parallel_rank: int | None = None, + reasoning_ended: bool | None = None, *, resumable: bool = False, ) -> None: @@ -1024,6 +1056,12 @@ class AsyncOmniEngine: sampling_params_list=sampling_params_list, final_stage_id=final_stage_id, arrival_time=arrival_time, + lora_request=lora_request, + tokenization_kwargs=tokenization_kwargs, + trace_headers=trace_headers, + priority=priority, + data_parallel_rank=data_parallel_rank, + reasoning_ended=reasoning_ended, resumable=resumable, ) diff --git a/vllm_omni/engine/stage_init_utils.py b/vllm_omni/engine/stage_init_utils.py index 6e8137206..f71afad83 100644 --- a/vllm_omni/engine/stage_init_utils.py +++ b/vllm_omni/engine/stage_init_utils.py @@ -467,6 +467,35 @@ def initialize_diffusion_stage( return StageDiffusionClient(model, od_config, metadata, batch_size=batch_size) +def _shutdown_or_close_resource(resource: Any, resource_name: str, stage_id: int) -> None: + """vLLM CoreEngineProcManager / coordinators use ``shutdown()``, not ``close()``.""" + if resource is None: + return + shutdown = getattr(resource, "shutdown", None) + if callable(shutdown): + try: + shutdown() + except Exception as cleanup_error: + logger.warning( + "[stage_init] Failed to shutdown launched %s for stage %s: %s", + resource_name, + stage_id, + cleanup_error, + ) + return + close = getattr(resource, "close", None) + if callable(close): + try: + close() + except Exception as cleanup_error: + logger.warning( + "[stage_init] Failed to close launched %s for stage %s: %s", + resource_name, + stage_id, + cleanup_error, + ) + + def close_started_llm_stage(started: StartedLlmStage) -> None: """Terminate the subprocess owned by a launched stage that never attached.""" if started.proc is None: diff --git a/vllm_omni/entrypoints/async_omni.py b/vllm_omni/entrypoints/async_omni.py index 6c8022461..129ef3c99 100644 --- a/vllm_omni/entrypoints/async_omni.py +++ b/vllm_omni/entrypoints/async_omni.py @@ -9,7 +9,7 @@ from __future__ import annotations import asyncio import time -from collections.abc import AsyncGenerator, Iterable, Sequence +from collections.abc import AsyncGenerator, Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any from vllm import TokensPrompt @@ -88,7 +88,12 @@ class AsyncOmni(EngineClient, OmniBase): else: vllm_config = self.engine.stage_vllm_configs[stage_index] io_processor_plugin = vllm_config.model_config.io_processor_plugin - self.io_processor = get_io_processor(vllm_config, io_processor_plugin) + renderer = self.renderer + if renderer is None: + from vllm.renderers import renderer_from_config + + renderer = renderer_from_config(vllm_config) + self.io_processor = get_io_processor(vllm_config, renderer, io_processor_plugin) def _get_comprehension_stage_index(self) -> int | None: fallback_idx: int | None = None @@ -159,6 +164,10 @@ class AsyncOmni(EngineClient, OmniBase): tokenization_kwargs: dict[str, Any] | None = None, sampling_params_list: Sequence[OmniSamplingParams] | None = None, output_modalities: list[str] | None = None, + trace_headers: Mapping[str, str] | None = None, + priority: int = 0, + data_parallel_rank: int | None = None, + reasoning_ended: bool | None = None, ) -> AsyncGenerator[OmniRequestOutput, None]: """Generate outputs for the given prompt(s) asynchronously. diff --git a/vllm_omni/entrypoints/cli/main.py b/vllm_omni/entrypoints/cli/main.py index 629a4641c..affa6c833 100644 --- a/vllm_omni/entrypoints/cli/main.py +++ b/vllm_omni/entrypoints/cli/main.py @@ -28,6 +28,10 @@ def main(): cli_env_setup() + from vllm_omni.entrypoints.cli.serve import _ensure_vllm_platform + + _ensure_vllm_platform() + parser = FlexibleArgumentParser( description="vLLM OMNI CLI", epilog=VLLM_SUBCMD_PARSER_EPILOG.format(subcmd="[subcommand]"), diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 4e1c8d3a9..b72df41cd 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -42,6 +42,37 @@ Search by using: `--help=` to explore options by section (e.g., """ +def _ensure_vllm_platform(): + """Ensure vLLM's current_platform is valid before arg parsing. + + Upstream vLLM's argument parser now instantiates DeviceConfig during + ``make_arg_parser``, which requires a resolved platform with a non-empty + ``device_type``. In some environments (e.g. editable installs with + broken package metadata), vLLM's own platform auto-detection may fail + and fall back to ``UnspecifiedPlatform``. When that happens, use the + Omni platform (which has its own detection logic) as a drop-in + replacement so that argument parsing succeeds. + """ + from vllm import platforms as vllm_platforms + + if vllm_platforms.current_platform.is_unspecified(): + from vllm_omni.platforms import current_omni_platform + + if not current_omni_platform.is_unspecified(): + vllm_platforms.current_platform = current_omni_platform + logger.debug( + "Replaced vLLM UnspecifiedPlatform with omni platform %s", + type(current_omni_platform).__name__, + ) + else: + from vllm.platforms.cpu import CpuPlatform + + vllm_platforms.current_platform = CpuPlatform() + logger.debug( + "Both vLLM and omni platforms are unspecified, falling back to CpuPlatform for arg parsing", + ) + + class OmniServeCommand(CLISubcommand): """The `serve` subcommand for the vLLM CLI.""" @@ -82,6 +113,7 @@ class OmniServeCommand(CLISubcommand): usage="vllm serve [model_tag] --omni [options]", ) + _ensure_vllm_platform() serve_parser = make_arg_parser(serve_parser) serve_parser.epilog = VLLM_SUBCMD_PARSER_EPILOG.format(subcmd=self.name) diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 4a7b097b2..627174b20 100644 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -639,6 +639,7 @@ async def omni_init_app_state( OpenAIServingResponses( engine_client, state.openai_serving_models, + openai_serving_render=state.openai_serving_render, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -699,7 +700,8 @@ async def omni_init_app_state( OpenAIServingPooling( engine_client, state.openai_serving_models, - supported_tasks=supported_tasks, + state.openai_serving_render, + supported_tasks=tuple(supported_tasks), request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -746,6 +748,7 @@ async def omni_init_app_state( state.openai_serving_tokenization = OpenAIServingTokenization( engine_client, state.openai_serving_models, + state.openai_serving_render, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -787,6 +790,7 @@ async def omni_init_app_state( 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, + default_chat_template_kwargs=args.default_chat_template_kwargs, ) if "generate" in supported_tasks else None @@ -795,6 +799,7 @@ async def omni_init_app_state( ServingTokens( engine_client, state.openai_serving_models, + state.openai_serving_render, request_logger=request_logger, return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_prompt_tokens_details=args.enable_prompt_tokens_details, diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 35f56516c..e84a49aac 100644 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -62,7 +62,7 @@ from vllm.entrypoints.openai.parser.harmony_utils import ( from vllm.entrypoints.openai.responses.protocol import ResponsesRequest 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 import PromptType from vllm.logger import init_logger from vllm.outputs import RequestOutput from vllm.reasoning import ReasoningParser @@ -818,7 +818,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Prepare the tool parser if it's needed try: if tool_choice_auto and self.tool_parser: - tool_parsers: list[ToolParser | None] = [self.tool_parser(tokenizer)] * num_choices + tool_parsers: list[ToolParser | None] = [self.tool_parser(tokenizer, request.tools)] * num_choices else: tool_parsers = [None] * num_choices except Exception as e: @@ -1639,12 +1639,12 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): logprobs = None if self.use_harmony: - reasoning_content, content, _ = parse_chat_output(token_ids) + reasoning, content, _ = parse_chat_output(token_ids) if not request.include_reasoning: - reasoning_content = None + reasoning = None if self.tool_parser is not None: - tool_parser = self.tool_parser(tokenizer) + tool_parser = self.tool_parser(tokenizer, request.tools) # NOTE: We use token_ids for openai tool parser tool_call_info = tool_parser.extract_tool_calls( "", @@ -1654,14 +1654,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): content = tool_call_info.content message = ChatMessage( role=role, - reasoning_content=reasoning_content, + reasoning=reasoning, content=content, tool_calls=tool_call_info.tool_calls, ) else: message = ChatMessage( role=role, - reasoning_content=reasoning_content, + reasoning=reasoning, content=content, ) @@ -1682,11 +1682,11 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if reasoning_parser: # If the reasoning parser is enabled, # tool calls are extracted exclusively from the content. - reasoning_content, content = reasoning_parser.extract_reasoning(output.text, request=request) + reasoning, content = reasoning_parser.extract_reasoning(output.text, request=request) if not request.include_reasoning: - reasoning_content = None + reasoning = None else: - reasoning_content = None + reasoning = None content = output.text auto_tools_called = False @@ -1696,14 +1696,14 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): not isinstance(request.tool_choice, ChatCompletionNamedToolChoiceParam) and request.tool_choice != "required" ): - message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content) + message = ChatMessage(role=role, reasoning=reasoning, content=content) # if the request uses tools and specified a tool choice elif request.tool_choice and type(request.tool_choice) is ChatCompletionNamedToolChoiceParam: tool_call_class = MistralToolCall if isinstance(tokenizer, MistralTokenizer) else ToolCall message = ChatMessage( role=role, - reasoning_content=reasoning_content, + reasoning=reasoning, content="", tool_calls=[ tool_call_class( @@ -1745,13 +1745,13 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) for i, tool_call in enumerate(tool_calls) ], - reasoning_content=reasoning_content, + reasoning=reasoning, ) # if the request doesn't use tool choice # OR specifies to not use a tool elif not request.tool_choice or request.tool_choice == "none": - message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content) + message = ChatMessage(role=role, reasoning=reasoning, content=content) # handle when there are tools and tool choice is auto elif ( @@ -1761,7 +1761,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): and self.tool_parser ): try: - tool_parser = self.tool_parser(tokenizer) + tool_parser = self.tool_parser(tokenizer, request.tools) except RuntimeError as e: logger.exception("Error in tool parser creation.") return self.create_error_response(e) @@ -1774,7 +1774,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if tool_call_info.tools_called: message = ChatMessage( role=role, - reasoning_content=reasoning_content, + reasoning=reasoning, content=tool_call_info.content, tool_calls=tool_call_info.tool_calls, ) @@ -1790,7 +1790,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ret_content = tool_call_info.content message = ChatMessage( role=role, - reasoning_content=reasoning_content, + reasoning=reasoning, content=ret_content, ) @@ -1800,7 +1800,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): "Error in chat_completion_full_generator - cannot determine if tools should be extracted. " "Returning a standard chat completion." ) - message = ChatMessage(role=role, reasoning_content=reasoning_content, content=content) + message = ChatMessage(role=role, reasoning=reasoning, content=content) choice_data = ChatCompletionResponseChoice( index=output.index, diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 8126fd544..f05126882 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -1642,3 +1642,6 @@ class OmniOpenAIServingSpeech(OpenAIServing, AudioMixin): succeeded=succeeded, failed=len(final_results) - succeeded, ) + + +ServingSpeech = OmniOpenAIServingSpeech diff --git a/vllm_omni/entrypoints/utils.py b/vllm_omni/entrypoints/utils.py index c664fe80a..e29e9eea1 100644 --- a/vllm_omni/entrypoints/utils.py +++ b/vllm_omni/entrypoints/utils.py @@ -1,7 +1,7 @@ import os import types from collections import Counter -from dataclasses import asdict, fields, is_dataclass +from dataclasses import fields, is_dataclass from pathlib import Path from typing import Any, get_args, get_origin @@ -145,12 +145,21 @@ def _convert_dataclasses_to_dict(obj: Any) -> Any: if isinstance(obj, set): return list(obj) # Handle dataclass objects - # Note: asdict() recursively converts nested dataclasses but not Counter objects, - # so we need to recursively process the result - if is_dataclass(obj): - result = asdict(obj) - # Recursively process the result to convert any Counter objects - return _convert_dataclasses_to_dict(result) + # Use field iteration instead of asdict() to: + # 1. Only include init fields (non-init fields cause "unexpected kwarg" errors) + # 2. Skip None values matching field defaults (avoids Pydantic validation + # failures when None is explicitly passed for non-Optional typed fields, + # e.g. CompilationConfig.cudagraph_capture_sizes: list[int] = None) + if is_dataclass(obj) and not isinstance(obj, type): + result = {} + for f in fields(obj): + if not f.init: + continue + value = getattr(obj, f.name) + if value is None and f.default is None: + continue + result[f.name] = _convert_dataclasses_to_dict(value) + return result # Handle dictionaries (recurse into values) and filter out callables(cause error in OmegaConf.create) # Note: This must come AFTER Counter check since Counter is a dict subclass if isinstance(obj, dict): diff --git a/vllm_omni/inputs/data.py b/vllm_omni/inputs/data.py index 5768c3b6d..7824e7092 100644 --- a/vllm_omni/inputs/data.py +++ b/vllm_omni/inputs/data.py @@ -16,7 +16,7 @@ except ImportError: import torch -from vllm.inputs.data import EmbedsPrompt, TextPrompt, TokenInputs, TokensPrompt +from vllm.inputs import EmbedsPrompt, TextPrompt, TokensInput, TokensPrompt class OmniTextPrompt(TextPrompt): @@ -59,10 +59,10 @@ class OmniTokensPrompt(TokensPrompt): additional_information: NotRequired[dict[str, Any]] -class OmniTokenInputs(TokenInputs): +class OmniTokenInputs(TokensInput): """Token inputs with optional embeddings and additional information. - Extends TokenInputs to support prompt embeddings and additional + Extends TokensInput to support prompt embeddings and additional information payloads for direct transfer between pipeline stages. Attributes: diff --git a/vllm_omni/inputs/preprocess.py b/vllm_omni/inputs/preprocess.py index 15f31627f..c6dffd054 100644 --- a/vllm_omni/inputs/preprocess.py +++ b/vllm_omni/inputs/preprocess.py @@ -1,10 +1,9 @@ from typing import Any from typing_extensions import assert_never -from vllm.inputs.data import EmbedsInputs, SingletonInputs +from vllm.inputs import EmbedsInput, MultiModalInput, SingletonInput from vllm.inputs.preprocess import InputPreprocessor from vllm.logger import init_logger -from vllm.multimodal.inputs import MultiModalInputs from vllm.renderers.inputs import SingletonDictPrompt from vllm_omni.inputs.data import ( @@ -30,7 +29,7 @@ class OmniInputPreprocessor(InputPreprocessor): self, parsed_content: OmniTextPrompt, tokenization_kwargs: dict[str, Any] | None = None, - ) -> OmniTokenInputs | MultiModalInputs: + ) -> OmniTokenInputs | MultiModalInput: """Process text prompts with support for mm_processor_kwargs. Extends base class to support mm_processor_kwargs without multi_modal_data. @@ -40,7 +39,7 @@ class OmniInputPreprocessor(InputPreprocessor): prompt_text = parsed_content["prompt"] mm_processor_kwargs = parsed_content.get("mm_processor_kwargs") or {} - inputs: OmniTokenInputs | MultiModalInputs + inputs: OmniTokenInputs | MultiModalInput if multi_modal_data := parsed_content.get("multi_modal_data"): inputs = self._process_multimodal( prompt_text, @@ -86,14 +85,14 @@ class OmniInputPreprocessor(InputPreprocessor): self, parsed_content: OmniTokensPrompt, tokenization_kwargs: dict[str, Any] | None = None, - ) -> OmniTokenInputs | MultiModalInputs: + ) -> OmniTokenInputs | MultiModalInput: prompt_token_ids = self._truncate_inputs(parsed_content["prompt_token_ids"], tokenization_kwargs) prompt_embeds = parsed_content.get("prompt_embeds") additional_information = parsed_content.get("additional_information") multi_modal_data = parsed_content.get("multi_modal_data") - inputs: OmniTokenInputs | MultiModalInputs + inputs: OmniTokenInputs | MultiModalInput if multi_modal_data: inputs = self._process_multimodal( prompt_token_ids, @@ -123,7 +122,7 @@ class OmniInputPreprocessor(InputPreprocessor): def _process_embeds( self, parsed_content: OmniEmbedsPrompt, - ) -> EmbedsInputs: + ) -> EmbedsInput: """Process embeddings prompt with omni-specific extensions. Extends base _process_embeds to handle additional_information payload @@ -143,7 +142,7 @@ class OmniInputPreprocessor(InputPreprocessor): self, prompt: SingletonDictPrompt, tokenization_kwargs: dict[str, Any] | None = None, - ) -> SingletonInputs: + ) -> SingletonInput: """ Extract the singleton inputs from a prompt. @@ -153,7 +152,7 @@ class OmniInputPreprocessor(InputPreprocessor): Returns: - * [`SingletonInputs`][vllm.inputs.data.SingletonInputs] instance + * [`SingletonInput`][vllm.inputs.engine.SingletonInput] instance """ if "prompt_embeds" in prompt: return self._process_embeds(prompt) # type: ignore[arg-type] diff --git a/vllm_omni/model_executor/models/bagel/bagel.py b/vllm_omni/model_executor/models/bagel/bagel.py index e79f0212e..934f434e6 100644 --- a/vllm_omni/model_executor/models/bagel/bagel.py +++ b/vllm_omni/model_executor/models/bagel/bagel.py @@ -8,6 +8,7 @@ import torch.nn as nn from transformers import BatchFeature from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions +from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.layernorm import RMSNorm as VllmRMSNorm from vllm.model_executor.layers.linear import ( QKVParallelLinear, @@ -20,7 +21,6 @@ from vllm.model_executor.models.qwen2 import Qwen2DecoderLayer, Qwen2MLP from vllm.model_executor.models.utils import AutoWeightsLoader from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( - MultiModalDataDict, MultiModalFieldConfig, MultiModalKwargsItems, ) diff --git a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py index 784393e18..bc04aae33 100644 --- a/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py +++ b/vllm_omni/model_executor/models/cosyvoice3/cosyvoice3.py @@ -10,10 +10,11 @@ import torch.nn as nn from transformers.feature_extraction_utils import BatchFeature from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions +from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.models.interfaces import SupportsMultiModal from vllm.multimodal import MULTIMODAL_REGISTRY -from vllm.multimodal.inputs import MultiModalDataDict, MultiModalFieldConfig, MultiModalKwargsItems +from vllm.multimodal.inputs import MultiModalFieldConfig, MultiModalKwargsItems from vllm.multimodal.parse import MultiModalDataItems, MultiModalDataParser from vllm.multimodal.processing import ( BaseDummyInputsBuilder, diff --git a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py index f90826fd3..31eed9b2c 100644 --- a/vllm_omni/model_executor/models/glm_image/glm_image_ar.py +++ b/vllm_omni/model_executor/models/glm_image/glm_image_ar.py @@ -40,6 +40,7 @@ from vllm.config import CacheConfig, MultiModalConfig, VllmConfig from vllm.config.multimodal import BaseDummyOptions from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size from vllm.distributed import utils as dist_utils +from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.attention import Attention from vllm.model_executor.layers.attention.mm_encoder_attention import ( @@ -73,7 +74,6 @@ from vllm.model_executor.models.utils import ( from vllm.model_executor.models.vision import get_vit_attn_backend from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( - MultiModalDataDict, MultiModalFeatureSpec, MultiModalFieldConfig, MultiModalKwargsItems, diff --git a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py index 5a0ae9965..6d25274f9 100644 --- a/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py +++ b/vllm_omni/model_executor/models/hunyuan_image3/hunyuan_image3.py @@ -20,6 +20,7 @@ from vllm.compilation.decorators import support_torch_compile from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions from vllm.distributed import get_pp_group +from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.fused_moe import SharedFusedMoE from vllm.model_executor.layers.linear import ( @@ -58,7 +59,6 @@ from vllm.model_executor.models.utils import ( from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.image import rgba_to_rgb from vllm.multimodal.inputs import ( - MultiModalDataDict, MultiModalFeatureSpec, MultiModalFieldConfig, MultiModalKwargsItems, diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py index 9acb81bce..22a9a9111 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio.py @@ -10,6 +10,7 @@ from torch import nn from transformers import BatchFeature, Qwen2Config from vllm.config import VllmConfig from vllm.config.multimodal import BaseDummyOptions +from vllm.inputs import ModalityData, MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.models import SupportsPP @@ -18,9 +19,7 @@ from vllm.model_executor.models.utils import init_vllm_registered_model from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( AudioItem, - ModalityData, MultiModalBatchedField, - MultiModalDataDict, MultiModalFieldConfig, MultiModalFieldElem, MultiModalKwargsItem, diff --git a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py index 1424ca775..56cb8788e 100644 --- a/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py +++ b/vllm_omni/model_executor/models/mimo_audio/mimo_audio_llm.py @@ -13,6 +13,7 @@ from transformers.models.qwen2.modeling_qwen2 import ( ) from vllm.config import VllmConfig from vllm.forward_context import get_forward_context +from vllm.inputs import MultiModalDataDict from vllm.model_executor.layers.linear import ColumnParallelLinear from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.model_loader.weight_utils import default_weight_loader, maybe_remap_kv_scale_name @@ -34,7 +35,6 @@ from vllm.model_executor.models.utils import ( ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( - MultiModalDataDict, MultiModalFieldConfig, MultiModalKwargsItems, ) 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 04212ceeb..ed6df6af3 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_omni.py @@ -18,7 +18,7 @@ from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import ( Qwen3OmniMoeThinkerConfig, ) from vllm.config import ModelConfig, VllmConfig -from vllm.inputs.data import PromptType, TokensPrompt +from vllm.inputs import PromptType, TokensPrompt from vllm.logger import init_logger from vllm.model_executor.layers.rotary_embedding import MRotaryEmbedding from vllm.model_executor.models.interfaces import SupportsMRoPE, SupportsMultiModal, SupportsPP, SupportsRealtime 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 fd7b14ab4..671ffb6cb 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 @@ -46,7 +46,7 @@ from transformers.models.whisper import WhisperFeatureExtractor from vllm.compilation.decorators import support_torch_compile from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig from vllm.distributed import get_pp_group, get_tensor_model_parallel_world_size -from vllm.inputs.data import PromptType +from vllm.inputs import PromptType from vllm.logger import init_logger from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY from vllm.model_executor.layers.attention.mm_encoder_attention import ( diff --git a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py index 3b8927d63..b5d116173 100644 --- a/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py +++ b/vllm_omni/model_executor/models/voxtral_tts/voxtral_tts_audio_generation.py @@ -29,6 +29,7 @@ from mistral_common.tokens.tokenizers.audio import Audio, AudioEncoder from transformers import BatchFeature from transformers.tokenization_utils_base import TextInput from vllm.config import VllmConfig +from vllm.inputs import MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.model_loader.weight_utils import default_weight_loader from vllm.model_executor.models.interfaces import SupportsMultiModal @@ -39,7 +40,6 @@ from vllm.model_executor.models.utils import ( ) from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( - MultiModalDataDict, MultiModalFieldConfig, MultiModalKwargsItems, NestedTensors, diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 6daef1a23..eafff821a 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -1,7 +1,7 @@ import sys from aenum import extend_enum -from vllm.inputs.data import TokensPrompt as _OriginalTokensPrompt +from vllm.inputs import TokensPrompt as _OriginalTokensPrompt from vllm.model_executor.layers.rotary_embedding import ( MRotaryEmbedding as _OriginalMRotaryEmbedding, ) diff --git a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py index 3d9cb86ba..138948064 100644 --- a/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py +++ b/vllm_omni/platforms/npu/worker/npu_ar_model_runner.py @@ -26,8 +26,8 @@ from vllm.v1.outputs import ( from vllm.v1.spec_decode.metadata import SpecDecodeMetadata from vllm.v1.structured_output.utils import apply_grammar_bitmask from vllm.v1.utils import record_function_or_nullcontext -from vllm.v1.worker import mamba_utils from vllm.v1.worker.gpu_model_runner import AsyncGPUModelRunnerOutput, PerLayerAttnMetadata +from vllm.v1.worker.mamba_utils import preprocess_mamba from vllm.v1.worker.ubatch_utils import maybe_create_ubatch_slices from vllm_ascend.ascend_forward_context import set_ascend_forward_context from vllm_ascend.attention.utils import AscendCommonAttentionMetadata @@ -243,7 +243,7 @@ class NPUARModelRunner(OmniNPUModelRunner): # '_update_states_after_model_execute', which is not overridden in vLLM-Ascend. # We simply utilize the implementation in vLLM. if self.cache_config.mamba_cache_mode == "align": - mamba_utils.preprocess_mamba( + preprocess_mamba( scheduler_output, self.kv_cache_config, self.cache_config, diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index 155b75675..f1115ab4c 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -137,8 +137,10 @@ class GPUARModelRunner(OmniGPUModelRunner): 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) + if has_kv_transfer_group(): + kv_connector_metadata = scheduler_output.kv_connector_metadata + if kv_connector_metadata is not None: + get_kv_transfer_group().handle_preemptions(kv_connector_metadata) num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens with ( @@ -146,7 +148,7 @@ class GPUARModelRunner(OmniGPUModelRunner): self.synchronize_input_prep(), ): # Update persistent batch states. - self._update_states(scheduler_output) + deferred_state_corrections_fn = self._update_states(scheduler_output) if has_ec_transfer() and not get_ec_transfer().is_consumer: with self.maybe_get_ec_connector_output( @@ -417,6 +419,9 @@ class GPUARModelRunner(OmniGPUModelRunner): ) self.kv_connector_output = kv_connector_output + if deferred_state_corrections_fn: + deferred_state_corrections_fn() + return None @torch.inference_mode() @@ -474,8 +479,11 @@ class GPUARModelRunner(OmniGPUModelRunner): with record_function_or_nullcontext("gpu_model_runner: sample"): sampler_output = self._sample(logits, spec_decode_metadata) + self._update_states_after_model_execute(sampler_output.sampled_token_ids, scheduler_output) + self._draft_token_ids = None self._draft_token_req_ids = None + self.valid_sampled_token_count_gpu = None self.input_batch.prev_sampled_token_ids = None def propose_draft_token_ids(sampled_token_ids): @@ -515,7 +523,7 @@ class GPUARModelRunner(OmniGPUModelRunner): 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, + self.optimistic_seq_lens_cpu, sampled_token_ids, self.requests, self.input_batch, diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 4db683a8b..d95b676f6 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -94,8 +94,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): 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) + if has_kv_transfer_group(): + kv_connector_metadata = scheduler_output.kv_connector_metadata + if kv_connector_metadata is not None: + get_kv_transfer_group().handle_preemptions(kv_connector_metadata) num_scheduled_tokens = scheduler_output.total_num_scheduled_tokens with ( @@ -104,7 +106,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): ): if self.model_config.async_chunk and num_scheduled_tokens: self._update_request_states(scheduler_output) - self._update_states(scheduler_output) + deferred_state_corrections_fn = self._update_states(scheduler_output) if not scheduler_output.total_num_scheduled_tokens: return EMPTY_MODEL_RUNNER_OUTPUT @@ -309,6 +311,10 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): slot_mappings, # OMNI: pass slot_mappings for upstream v1 API compatibility ) self.kv_connector_output = kv_connector_output + + if deferred_state_corrections_fn: + deferred_state_corrections_fn() + return None @torch.inference_mode() @@ -636,11 +642,14 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] # type: ignore[assignment] else: seq_lens = max_query_len # type: ignore[assignment] - self.seq_lens.np[:num_reqs] = seq_lens - self.seq_lens.np[num_reqs:] = 0 - self.seq_lens.copy_to_gpu() + self.seq_lens[:num_reqs] = ( + seq_lens + if isinstance(seq_lens, int) + else torch.tensor(seq_lens, dtype=torch.int32, device=self.device) + ) + self.seq_lens[num_reqs:] = 0 - cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self._arange_scratch) self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens self.query_start_loc.copy_to_gpu() @@ -696,7 +705,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner): elif self.uses_xdrope_dim > 0: positions = self.xdrope_positions.gpu[:, :num_tokens_padded] else: - positions = self.positions.gpu[:num_tokens_padded] + positions = self.positions[:num_tokens_padded] if get_pp_group().is_first_rank: intermediate_tensors = None diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 8e5689986..a7abaf7b6 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -243,7 +243,7 @@ class OmniGPUModelRunner(GPUModelRunner): mrope_pos_ptr += completion_part_len - def _update_states(self, scheduler_output: "SchedulerOutput") -> None: + def _update_states(self, scheduler_output: "SchedulerOutput"): """Update the cached states and the persistent batch with the scheduler output. @@ -271,7 +271,7 @@ class OmniGPUModelRunner(GPUModelRunner): # Zero GPU memory for freshly allocated cache blocks to prevent # stale NaN/data from corrupting attention or SSM computation. - if hasattr(scheduler_output, "new_block_ids_to_zero") and scheduler_output.new_block_ids_to_zero: + if scheduler_output.new_block_ids_to_zero: self._zero_block_ids(scheduler_output.new_block_ids_to_zero) # Free the cached encoder outputs. @@ -300,7 +300,11 @@ class OmniGPUModelRunner(GPUModelRunner): for req_id in unscheduled_req_ids: self.input_batch.remove_request(req_id) + if self.use_async_spec_decode: + self.prev_num_draft_tokens.np.fill(0) + reqs_to_add: list[CachedRequestState] = [] + deferred_spec_decode_corrections = [] # Add new requests to the cached states. for new_req_data in scheduler_output.scheduled_new_reqs: req_id = new_req_data.req_id @@ -398,10 +402,6 @@ class OmniGPUModelRunner(GPUModelRunner): 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] @@ -427,12 +427,18 @@ class OmniGPUModelRunner(GPUModelRunner): 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) + optimistic_num_accepted = req_state.prev_num_draft_len + req_state.output_token_ids.extend([-1] * optimistic_num_accepted) + + deferred_spec_decode_corrections.append((req_id, optimistic_num_accepted, req_state)) + + prev_req_index = ( + self.input_batch.prev_req_id_to_index.get(req_id) + if self.input_batch.prev_req_id_to_index + else None + ) + if prev_req_index is not None: + self.prev_num_draft_tokens.np[prev_req_index] = optimistic_num_accepted # Update the cached states. req_state.num_computed_tokens = num_computed_tokens @@ -449,7 +455,8 @@ class OmniGPUModelRunner(GPUModelRunner): 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. + # failure, or output_token_ids was inflated by the optimistic + # extend above (async spec decode). 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 @@ -513,6 +520,35 @@ class OmniGPUModelRunner(GPUModelRunner): # Refresh batch metadata with any pending updates. self.input_batch.refresh_metadata() + if deferred_spec_decode_corrections: + + def correct_spec_decode_token_counts(): + valid_sampled_token_count = self._get_valid_sampled_token_count() + if not valid_sampled_token_count: + return + prev_req_id_to_index = self.input_batch.prev_req_id_to_index + if not prev_req_id_to_index: + return + for ( + req_id, + optimistic_num_accepted, + req_state, + ) in deferred_spec_decode_corrections: + prev_req_index = prev_req_id_to_index.get(req_id) + if prev_req_index is None: + continue + num_accepted = valid_sampled_token_count[prev_req_index] - 1 + correction = optimistic_num_accepted - num_accepted + req_state.num_computed_tokens -= correction + cur_req_index = self.input_batch.req_id_to_index.get(req_id) + if cur_req_index is None: + continue + self.input_batch.num_computed_tokens_cpu[cur_req_index] -= correction + + return correct_spec_decode_token_counts + else: + return None + @torch.inference_mode() def extract_multimodal_outputs(self, hidden_states: torch.Tensor | list[torch.Tensor] | OmniOutput) -> dict: if ( @@ -704,11 +740,14 @@ class OmniGPUModelRunner(GPUModelRunner): seq_lens = [1] * num_decode_tokens + [num_prefill_tokens + 1] # type: ignore[assignment] else: seq_lens = max_query_len # type: ignore[assignment] - self.seq_lens.np[:num_reqs] = seq_lens - self.seq_lens.np[num_reqs:] = 0 - self.seq_lens.copy_to_gpu() + self.seq_lens[:num_reqs] = ( + seq_lens + if isinstance(seq_lens, int) + else torch.tensor(seq_lens, dtype=torch.int32, device=self.device) + ) + self.seq_lens[num_reqs:] = 0 - cum_num_tokens, _ = self._get_cumsum_and_arange(num_scheduled_tokens) + cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self._arange_scratch) self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens self.query_start_loc.copy_to_gpu() @@ -759,7 +798,7 @@ class OmniGPUModelRunner(GPUModelRunner): elif self.uses_xdrope_dim > 0: positions = self.xdrope_positions.gpu[:, :num_tokens_padded] else: - positions = self.positions.gpu[:num_tokens_padded] + positions = self.positions[:num_tokens_padded] if get_pp_group().is_first_rank: intermediate_tensors = None @@ -1164,7 +1203,7 @@ class OmniGPUModelRunner(GPUModelRunner): elif self.uses_xdrope_dim > 0: positions = self.xdrope_positions.gpu[:, :num_input_tokens] else: - positions = self.positions.gpu[:num_input_tokens] + positions = self.positions[:num_input_tokens] if is_first_rank: intermediate_tensors = None