From 7b837944a16ff440df0b19e71c6eca310d8dfc36 Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Sun, 14 Jun 2026 00:43:02 +0800 Subject: [PATCH] [Rebase] Rebase to vllm 0.23.0 (#4286) Signed-off-by: tzhouam --- .buildkite/pipeline-intel.yaml | 2 +- docker/Dockerfile.ci | 3 +- docker/Dockerfile.cuda | 2 +- docker/Dockerfile.rocm | 2 +- docker/Dockerfile.xpu | 2 +- docs/getting_started/installation/gpu.md | 2 +- .../installation/gpu/cuda.inc.md | 16 ++-- .../installation/gpu/rocm.inc.md | 14 +-- docs/getting_started/quickstart.md | 6 +- recipes/README.md | 1 - requirements/common.txt | 15 ++- tests/benchmarks/test_serve_cli.py | 56 +++++++++-- tests/diffusion/layers/test_norm.py | 19 +++- .../e2e/online_serving/test_higgs_audio_v3.py | 1 + vllm_omni/benchmarks/patch/patch.py | 3 + vllm_omni/core/sched/omni_ar_scheduler.py | 5 + .../hunyuan_image3/hunyuan_fused_moe.py | 96 ++++++++++++++++--- .../mooncake_transfer_engine_connector.py | 9 ++ vllm_omni/engine/async_omni_engine.py | 22 +++++ vllm_omni/engine/output_processor.py | 14 ++- vllm_omni/entrypoints/cli/benchmark/main.py | 2 +- vllm_omni/entrypoints/cli/main.py | 2 +- vllm_omni/entrypoints/cli/serve.py | 2 +- vllm_omni/entrypoints/openai/api_server.py | 20 ++-- vllm_omni/entrypoints/openai/serving_chat.py | 35 ++++++- .../minicpmo_4_5/minicpmo_4_5_omni_llm.py | 57 ++--------- .../moss_tts_nano/modeling_moss_tts_nano.py | 15 ++- .../models/qwen3_omni/qwen3_moe.py | 6 +- .../models/qwen3_tts/prompt_embeds_builder.py | 8 +- vllm_omni/model_executor/models/utils.py | 56 +++++++++++ vllm_omni/worker/gpu_ar_model_runner.py | 58 +++++++++++ vllm_omni/worker/gpu_ar_worker.py | 4 +- .../worker/gpu_generation_model_runner.py | 1 + vllm_omni/worker/gpu_generation_worker.py | 4 +- vllm_omni/worker/gpu_model_runner.py | 1 + vllm_omni/worker/memory_utils.py | 55 +++++++++++ 36 files changed, 481 insertions(+), 135 deletions(-) create mode 100644 vllm_omni/worker/memory_utils.py diff --git a/.buildkite/pipeline-intel.yaml b/.buildkite/pipeline-intel.yaml index d94d31bdd..cb1d76a1c 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.22.0" + VLLM_VERSION: "v0.23.0" priority: 100 timeout_in_minutes: 60 soft_fail: false diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 31e2529b1..dd55cd603 100644 --- a/docker/Dockerfile.ci +++ b/docker/Dockerfile.ci @@ -2,7 +2,7 @@ # The pinned VLLM_PRECOMPILED_WHEEL_COMMIT and VLLM_BASE_TAG below must # be updated on every rebase to match the new upstream vllm commit/tag. ARG VLLM_BASE_IMAGE=vllm/vllm-openai -ARG VLLM_BASE_TAG=v0.22.0 +ARG VLLM_BASE_TAG=v0.23.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} @@ -16,6 +16,7 @@ RUN apt-get update && \ RUN uv pip install --system ".[dev]" + RUN ln -sf /usr/bin/python3 /usr/bin/python ENTRYPOINT [] diff --git a/docker/Dockerfile.cuda b/docker/Dockerfile.cuda index 21098db20..fdea270f0 100644 --- a/docker/Dockerfile.cuda +++ b/docker/Dockerfile.cuda @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai:v0.22.0 +ARG BASE_IMAGE=vllm/vllm-openai:v0.23.0 FROM ${BASE_IMAGE} ARG COMMON_WORKDIR=/app diff --git a/docker/Dockerfile.rocm b/docker/Dockerfile.rocm index 5b0fa35cf..a8409b9a0 100644 --- a/docker/Dockerfile.rocm +++ b/docker/Dockerfile.rocm @@ -1,4 +1,4 @@ -ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.22.0 +ARG BASE_IMAGE=vllm/vllm-openai-rocm:v0.23.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 5ccdebf53..0dc703a98 100644 --- a/docker/Dockerfile.xpu +++ b/docker/Dockerfile.xpu @@ -88,7 +88,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.22.0 +ARG VLLM_VERSION=v0.23.0 RUN git clone -b ${VLLM_VERSION} https://github.com/vllm-project/vllm WORKDIR /workspace/vllm diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index b35cc69c5..59d03a3e4 100644 --- a/docs/getting_started/installation/gpu.md +++ b/docs/getting_started/installation/gpu.md @@ -34,7 +34,7 @@ vLLM-Omni is a Python library that supports the following GPU variants. The libr ### Pre-built wheels -Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, and 0.22.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). +Note: Pre-built wheels are currently available for vLLM-Omni 0.11.0rc1, 0.12.0rc1, 0.14.0rc1, 0.14.0, 0.16.0, 0.18.0, 0.20.0, 0.21.0, 0.22.0, and 0.23.0. If you need a newer unreleased revision, please [build from source](https://docs.vllm.ai/projects/vllm-omni/en/latest/getting_started/installation/gpu/#build-wheel-from-source). === "NVIDIA CUDA" diff --git a/docs/getting_started/installation/gpu/cuda.inc.md b/docs/getting_started/installation/gpu/cuda.inc.md index a8f293afd..9e394b758 100644 --- a/docs/getting_started/installation/gpu/cuda.inc.md +++ b/docs/getting_started/installation/gpu/cuda.inc.md @@ -20,7 +20,7 @@ Therefore, it is recommended to install vLLM and vLLM-Omni with a **fresh new** vLLM-Omni is built based on vLLM. Please install it with command below. ```bash -uv pip install vllm==0.22.0 --torch-backend=auto +uv pip install vllm==0.23.0 --torch-backend=auto ``` #### Installation of vLLM-Omni @@ -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.22.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.23.0 release version of the library ```bash -uv pip install vllm==0.22.0 --torch-backend=auto +uv pip install vllm==0.23.0 --torch-backend=auto ``` -The 0.22.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.23.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,12 +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.22.0 +git checkout v0.23.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 -#For CUDA 12.9 (the default for v0.22.0; the wheel filename has no `+cu129` suffix) -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.22.0/vllm-0.22.0-cp38-abi3-manylinux_2_28_x86_64.whl +#For CUDA 13.0 (the default for v0.23.0; the wheel filename has no `+cu130` suffix) +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.23.0/vllm-0.23.0-cp38-abi3-manylinux_2_28_x86_64.whl ``` Install vllm with command below (If you have no existing PyTorch). ```bash @@ -102,7 +102,7 @@ docker run --runtime nvidia --gpus 2 \ --env "HF_TOKEN=$HF_TOKEN" \ -p 8091:8091 \ --ipc=host \ - vllm/vllm-omni:v0.22.0 \ + vllm/vllm-omni:v0.23.0 \ vllm serve Qwen/Qwen3-Omni-30B-A3B-Instruct --omni --port 8091 ``` diff --git a/docs/getting_started/installation/gpu/rocm.inc.md b/docs/getting_started/installation/gpu/rocm.inc.md index 3912edb0e..a8344b21b 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.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722 +uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722 ``` #### Installation of vLLM-Omni @@ -34,13 +34,13 @@ uv pip install onnxruntime-rocm # --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.22.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.23.0 release version of the library ```bash -uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722 +uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722 ``` -The pre-built 0.22.0 vLLM wheel targets ROCm 7.2.2. If you need a different ROCm stack or want to reuse an existing PyTorch installation, build vLLM from source instead. +The pre-built 0.23.0 vLLM wheel targets ROCm 7.2.2. 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.22.0 +git checkout v0.23.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.22.0 \ + vllm/vllm-omni-rocm:v0.23.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.22.0 + vllm/vllm-omni-rocm:v0.23.0 ``` # --8<-- [end:pre-built-images] diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 005c3f9e9..652c512a2 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.22.0 --torch-backend=auto +uv pip install vllm==0.23.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.22.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.22.0/rocm722 +uv pip install vllm==0.23.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.23.0/rocm722 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -35,7 +35,7 @@ For additional installation methods — please see the [installation guide](inst !!! note It is important to install the same major & minor version of vLLM and vLLM Omni, otherwise things may not work as expected. If the versions are misaligned, you will see a warning when you import vLLM Omni. - If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.22.0` and vLLM Omni `0.22.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. + If you are seeing strange behavior with the `vllm` command not handling the `--omni` flag correctly, you most likely have a version mismatch with vLLM < `0.23.0` and vLLM Omni `0.23.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. ## Offline Inference diff --git a/recipes/README.md b/recipes/README.md index 2bbc518df..94cbcdbd5 100644 --- a/recipes/README.md +++ b/recipes/README.md @@ -54,7 +54,6 @@ recipes/ | [`StabilityAI/Stable-Diffusion-3.5.md`](./StabilityAI/Stable-Diffusion-3.5.md) | Text-to-image serving (SD 3.5-medium and SD 3.5-large) | 1x RTX A6000 48GB | | [`zai-org/GLM-TTS.md`](./zai-org/GLM-TTS.md) | Online serving for Chinese/English zero-shot voice-cloned TTS | 1x A40 48GB | | [`GLM/GLM-Image.md`](./GLM/GLM-Image.md) | Online serving for image generation | 1x A800 80GB / 2x A800 80GB | -| Within a single recipe file, include different hardware support sections such as `GPU`, `ROCm`, and `NPU`, and add concrete tested configurations like diff --git a/requirements/common.txt b/requirements/common.txt index fb435218a..5cfaeef55 100644 --- a/requirements/common.txt +++ b/requirements/common.txt @@ -1,10 +1,15 @@ # Common dependencies for all platforms -# transformers: keep aligned with upstream vLLM constraint. +# transformers: kept exactly aligned with the upstream vLLM constraint. # vLLM pins ">= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0". -# 5.9.0 breaks models that define _keys_to_ignore_on_load_unexpected as a list -# (e.g. OpenMOSS-Team/MOSS-TTS-Nano) because it changed -# _adjust_missing_and_unexpected_keys from list+list to set|set. -transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0, < 5.9.0 +# The previous extra "< 5.9.0" exclusion was removed when re-aligning with +# upstream: transformers 5.9 changed _adjust_missing_and_unexpected_keys from +# list+list to set|set, breaking models that declare +# _keys_to_ignore_on_load_unexpected as a list (e.g. OpenMOSS-Team/MOSS-TTS-Nano) +# with "TypeError: unsupported operand type(s) for |: 'list' and 'set'". +# That breakage is now handled in code instead of by pinning: +# MossTTSNanoForGeneration wraps weight loading in +# transformers_keys_to_ignore_compat (vllm_omni/model_executor/models/utils.py). +transformers >= 4.56.0, != 5.0.*, != 5.1.*, != 5.2.*, != 5.3.*, != 5.4.*, != 5.5.0 av>=14.0.0 omegaconf>=2.3.0 diffusers==0.38.0 diff --git a/tests/benchmarks/test_serve_cli.py b/tests/benchmarks/test_serve_cli.py index ab7b09f7b..6970f6e2f 100644 --- a/tests/benchmarks/test_serve_cli.py +++ b/tests/benchmarks/test_serve_cli.py @@ -42,10 +42,12 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path): yield chunk class MockResponse: - def __init__(self): - self.status = 200 - self.reason = "OK" - self.content = _Content(SSE_CHUNKS) + def __init__(self, status=200, reason="OK", chunks=None): + self.status = status + self.reason = reason + self.content = _Content(chunks or SSE_CHUNKS) + self._json_data = None + self._text_data = None async def __aenter__(self): return self @@ -53,20 +55,49 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path): async def __aexit__(self, exc_type, exc, tb): return False + async def text(self): + if self._text_data is not None: + return self._text_data + return json.dumps({"tokens": [1, 2, 3, 4, 5]}) + + async def json(self): + if self._json_data is not None: + return self._json_data + return {"tokens": [1, 2, 3, 4, 5]} + + def raise_for_status(self): + pass + class MockClientSession: def __init__(self, *args, **kwargs): pass + async def __aenter__(self): + return self + + async def __aexit__(self, exc_type, exc, tb): + return False + def post(self, url=None, *args, **kwargs): if url is not None: POST_CALLS.append(url) - return MockResponse() + return MockResponse(status=200) + + async def get(self, url=None, *args, **kwargs): + if url is not None: + POST_CALLS.append(url) + return MockResponse(status=200) async def close(self): return None - import vllm_omni.benchmarks.patch.patch as patch_mod + # Patch globally so modules loaded after this also see the mock + import aiohttp + aiohttp.ClientSession = MockClientSession + aiohttp.TCPConnector = lambda *args, **kwargs: object() + # Also patch in the patch module + import vllm_omni.benchmarks.patch.patch as patch_mod patch_mod.aiohttp.ClientSession = MockClientSession patch_mod.aiohttp.TCPConnector = lambda *args, **kwargs: object() @@ -136,8 +167,13 @@ def test_bench_serve_cli_mocks_http_request(tmp_path: Path): expected_url = f"http://127.0.0.1:{port}/v1/chat/completions" requested_urls = calls["requested_urls"] - sent_requests = len(requested_urls) + # Count only benchmark requests (to the chat completions endpoint), + # excluding tokenize/detokenize alignment calls from the upstream. + bench_requests = [url for url in requested_urls if url == expected_url] + sent_bench_requests = len(bench_requests) - assert result["completed"] == sent_requests == num_prompts - assert requested_urls - assert all(url == expected_url for url in requested_urls), f"Unexpected target URLs: {requested_urls}" + assert result["completed"] == sent_bench_requests == num_prompts, ( + f"completed={result['completed']}, bench_requests={sent_bench_requests}, all_requested_urls={requested_urls}" + ) + assert bench_requests + assert all(url == expected_url for url in bench_requests), f"Unexpected target URLs: {bench_requests}" diff --git a/tests/diffusion/layers/test_norm.py b/tests/diffusion/layers/test_norm.py index 3af50fa87..694f8f8e9 100644 --- a/tests/diffusion/layers/test_norm.py +++ b/tests/diffusion/layers/test_norm.py @@ -578,12 +578,27 @@ def test_rmsnorm_fused_native_parity_2d_input(): def test_rmsnorm_fused_native_parity_custom_weight(): - """Verify fused and native RMSNorm parity with non-default weights.""" + """Verify fused and native RMSNorm parity with non-default weights. + + Note: bfloat16 requires wider tolerance because the fused CUDA kernel + (vllm._custom_ops.rms_norm) computes in bf16 precision internally, + while forward_native uses float32 intermediates. With random custom + weights (not the default ones-weight), this precision gap exceeds 1e-3. + + The gap surfaced after upstream vLLM routed RMSNorm through the vLLM IR + op (ir.ops.rms_norm) — see upstream commit 40bb17502 + "[vLLM IR] 1/N Implement IR skeleton and rms_norm op (#33825)" — which + changed which fused kernel/precision path the layer dispatches to. + """ from vllm_omni.diffusion.layers.norm import RMSNorm hidden_size = 64 eps = 1e-6 + # Per-dtype tolerances: bf16 fused kernel has inherent precision limit + atol_map = {torch.float32: 1e-3, torch.float16: 1e-3, torch.bfloat16: 2e-2} + rtol_map = {torch.float32: 1e-3, torch.float16: 1e-3, torch.bfloat16: 1e-2} + for dtype in [torch.float32, torch.float16, torch.bfloat16]: torch.manual_seed(789) norm = RMSNorm(hidden_size, eps=eps).cuda() @@ -594,4 +609,4 @@ def test_rmsnorm_fused_native_parity_custom_weight(): out_fused = norm._forward_fused(x) out_native = norm.forward_native(x) - torch.testing.assert_close(out_fused, out_native, atol=1e-3, rtol=1e-3) + torch.testing.assert_close(out_fused, out_native, atol=atol_map[dtype], rtol=rtol_map[dtype]) diff --git a/tests/e2e/online_serving/test_higgs_audio_v3.py b/tests/e2e/online_serving/test_higgs_audio_v3.py index aaa5498b3..d6aecce41 100644 --- a/tests/e2e/online_serving/test_higgs_audio_v3.py +++ b/tests/e2e/online_serving/test_higgs_audio_v3.py @@ -140,6 +140,7 @@ class TestHiggsAudioV3OnlineHappyPath: } ) + @pytest.mark.skip(reason="issue#4411") @pytest.mark.core_model @pytest.mark.tts @hardware_test(res={"cuda": "H100"}, num_cards=1) diff --git a/vllm_omni/benchmarks/patch/patch.py b/vllm_omni/benchmarks/patch/patch.py index 76b8628c1..73720f4bd 100644 --- a/vllm_omni/benchmarks/patch/patch.py +++ b/vllm_omni/benchmarks/patch/patch.py @@ -1193,6 +1193,7 @@ async def benchmark( ramp_up_end_rps: int | None = None, ready_check_timeout_sec: int = 600, ssl_context: ssl.SSLContext | bool | None = None, + self_timed: bool = False, ): try: request_func = ASYNC_REQUEST_FUNCS[endpoint_type] @@ -1355,6 +1356,7 @@ async def benchmark( ramp_up_strategy, ramp_up_start_rps, ramp_up_end_rps, + self_timed, ): if ramp_up_strategy is not None: current_int_rps = int(current_request_rate) @@ -1443,6 +1445,7 @@ async def benchmark( defs.AVERAGE_PIXELS_PER_IMAGE: getattr(metrics, defs.AVERAGE_PIXELS_PER_IMAGE), defs.MEAN_DENOISE_STEP_LATENCY_MS: getattr(metrics, defs.MEAN_DENOISE_STEP_LATENCY_MS), "input_lens": [output.prompt_len for output in outputs], + "start_times": [output.start_time for output in outputs], "output_lens": actual_output_lens, "ttfts": [output.ttft for output in outputs], "itls": [output.itl for output in outputs], diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 736ba549c..984844bdb 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -712,6 +712,11 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): self._omits_kv_transfer_cache.pop(request.request_id, None) + # [Upstream compat] Discard request from in-flight prefills set added + # upstream for routed-experts in-flight reservation tracking. + # Use getattr for safety with test __new__ code paths. + getattr(self, "_inflight_prefills", set()).discard(request) + # 1. Standard cleanup parts from base _free_request connector_delay_free_blocks, kv_xfer_params = self._connector_finished(request) diff --git a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py index fe204502b..f9f114513 100644 --- a/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py +++ b/vllm_omni/diffusion/models/hunyuan_image3/hunyuan_fused_moe.py @@ -4,7 +4,6 @@ from typing import Any import vllm.forward_context as _vllm_fc -from vllm.model_executor.layers.fused_moe import FusedMoE from vllm.utils.import_utils import resolve_obj_by_qualname from vllm_omni.platforms import current_omni_platform @@ -25,22 +24,93 @@ def _set_forward_context_num_tokens(num_tokens: int) -> None: forward_context.in_profile_run = False -class HunyuanFusedMoEDefault(FusedMoE): - def __init__(self, *, prefix: str = "", **kwargs: Any) -> None: +class HunyuanFusedMoEDefault: + """Adapter that configures the upstream FusedMoE ``MoERunner`` for HunyuanImage3. + + Upstream commit dc68bd8c41 refactored FusedMoE from a class (``nn.Module``) + into a factory function that returns a ``MoERunner`` instance, whose expert + weights live in a ``routed_experts`` submodule + (``...experts.routed_experts.w13_weight`` / ``...w2_weight``). + + This adapter builds that runner, installs the omni-specific forward-context + setup and one-shot kernel-initialisation hook, and returns the runner + *directly* from ``__new__`` so the parent MoE block registers it as a real + ``nn.Module`` submodule. + + Returning the runner (rather than wrapping it in a plain object that holds it + in an attribute) is required for correctness: a non-Module wrapper hides the + runner's parameters from ``named_parameters()``, so ``load_weights`` cannot + find ``...experts.routed_experts.w13_weight`` and raises ``KeyError`` during + weight loading. + """ + + def __new__(cls, *, prefix: str = "", **kwargs: Any) -> Any: # Current vLLM FusedMoE handles output reduction internally. kwargs.pop("reduce_results", None) - super().__init__(prefix=prefix, **kwargs) - self._prefix = prefix - self._init_hook_handle = self.register_forward_pre_hook(self._initialize_kernel_hook, with_kwargs=True) + # FusedMoE is now a factory function — call it to get a MoERunner. + from vllm.model_executor.layers.fused_moe import FusedMoE as _FusedMoE - def _initialize_kernel_hook(self, module: Any, args: Any, kwargs: Any) -> None: - if self.quant_method and getattr(self.quant_method, "moe_kernel", None) is None: - self.quant_method.process_weights_after_loading(self) - self._init_hook_handle.remove() + moe_runner = _FusedMoE(prefix=prefix, **kwargs) - def forward(self, hidden_states: Any, router_logits: Any) -> Any: - _set_forward_context_num_tokens(hidden_states.shape[0]) - return super().forward(hidden_states, router_logits) + # Set ForwardContext.num_tokens before each forward. After the rebase to + # vLLM 0.18.0 FusedMoE requires this; without it MoE routing is silently + # incorrect. Previously done in the wrapper's forward(); now a pre-hook + # on the runner so we can return the runner directly. + def _num_tokens_pre_hook(module: Any, args: Any, kwargs: Any) -> None: + hidden_states = kwargs.get("hidden_states") + if hidden_states is None and args: + hidden_states = args[0] + if hidden_states is not None: + _set_forward_context_num_tokens(hidden_states.shape[0]) + + moe_runner.register_forward_pre_hook(_num_tokens_pre_hook, with_kwargs=True) + + # One-shot lazy kernel initialisation on the first forward (no-op unless + # the runner exposes an uninitialised quant_method). Mirrors the prior + # wrapper behaviour exactly, just bound to the runner module. + init_handle: Any = None + + def _kernel_init_pre_hook(module: Any, args: Any, kwargs: Any) -> None: + nonlocal init_handle + quant_method = getattr(module, "quant_method", None) + if quant_method is not None and getattr(quant_method, "moe_kernel", None) is None: + quant_method.process_weights_after_loading(module) + if init_handle is not None: + init_handle.remove() + + init_handle = moe_runner.register_forward_pre_hook(_kernel_init_pre_hook, with_kwargs=True) + + return moe_runner + + @staticmethod + def make_expert_params_mapping( + model: Any, + ckpt_gate_proj_name: str, + ckpt_down_proj_name: str, + ckpt_up_proj_name: str, + num_experts: int, + num_redundant_experts: int = 0, + ) -> list[tuple[str, str, int, str]]: + """Delegate to the upstream standalone function. + + Upstream vLLM refactored ``FusedMoE`` from a class (which had + ``make_expert_params_mapping`` as a classmethod) to a factory + function. The method was moved to a standalone function + ``fused_moe_make_expert_params_mapping`` in + ``vllm.model_executor.layers.fused_moe``. + """ + from vllm.model_executor.layers.fused_moe import ( + fused_moe_make_expert_params_mapping, + ) + + return fused_moe_make_expert_params_mapping( + model, + ckpt_gate_proj_name=ckpt_gate_proj_name, + ckpt_down_proj_name=ckpt_down_proj_name, + ckpt_up_proj_name=ckpt_up_proj_name, + num_experts=num_experts, + num_redundant_experts=num_redundant_experts, + ) class HunyuanFusedMoE: diff --git a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py index 1c058c7f0..297d4ef59 100644 --- a/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py +++ b/vllm_omni/distributed/omni_connectors/connectors/mooncake_transfer_engine_connector.py @@ -114,6 +114,7 @@ class MooncakeTransferEngineConnector(OmniConnectorBase): self._listener_ready = threading.Event() self._local_buffers: dict[str, Any] = {} self._local_buffers_lock = threading.Lock() + self._put_lock = threading.Lock() self._req_local = threading.local() self._worker_local = threading.local() self._last_ttl_check: float = _time_mod.monotonic() @@ -382,6 +383,14 @@ class MooncakeTransferEngineConnector(OmniConnectorBase): put_key = self._make_key(put_key, from_stage, to_stage) + # Serialize concurrent put() calls on the same connector to prevent + # races in the alloc -> copy -> store-metadata flow at the Mooncake + # C++ engine level. + with self._put_lock: + return self._put_impl(put_key, data) + + def _put_impl(self, put_key: str, data: Any) -> tuple[bool, int, dict[str, Any] | None]: + """Internal put implementation, called under _put_lock.""" try: src_addr = 0 size = 0 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index f53cb26ac..c654bda56 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -2480,6 +2480,28 @@ class AsyncOmniEngine: logger.exception("[AsyncOmniEngine] Failed to close OmniCoordinatorRuntime during shutdown") self._coordinator_runtime = None + # ── Release CuMem allocator memory pool ────────────────────────────── + # When enable_sleep_mode is in use, the CuMem (CUDA Virtual Memory + # Management) allocator holds model weights in a singleton memory pool + # that lives in the parent process. Killing the engine-core subprocess + # does NOT release this pool — the weights stay resident on the GPU + # and can cause CUDA OOM for subsequent engine instances (especially + # large models like BAGEL-7B-MoT whose weights alone consume ~134 GiB). + # + # Discard mode (level=2) is correct at shutdown: there is no benefit to + # keeping a CPU backup when the engine is being torn down. + try: + from vllm.device_allocator.cumem import CuMemAllocator, cumem_available + + if cumem_available: + allocator = CuMemAllocator.get_instance() + # Sleep at level 2 discards all pool memory from the GPU + # without creating CPU backups — cheapest and fastest. + allocator.sleep() + logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown") + except Exception: + pass + def _try_shutdown(self, *args, **kwargs) -> None: try: self.shutdown() diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index 5780c574c..b84a9eb3a 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -122,9 +122,11 @@ class OmniRequestState(RequestState): *args, **kwargs, ): + # arrival_time is always passed as a keyword argument by the sole caller + # (from_new_request). The *args fallback via positional index was removed + # because upstream RequestState.__init__ reorders parameters between + # releases, making args[12] fragile. arrival_time = kwargs.get("arrival_time") - if arrival_time is None and len(args) > 12: - arrival_time = args[12] super().__init__(*args, **kwargs) self.native_text_stats = RequestStateStats(arrival_time=float(arrival_time or 0.0)) # Omni-specific: multimodal output accumulation @@ -669,6 +671,14 @@ class MultimodalOutputProcessor(VLLMOutputProcessor): native_stats = req_state.native_text_stats if isinstance(req_state, OmniRequestState) else None previous_last_token_ts = native_stats.last_token_ts if native_stats is not None else 0.0 + # NOTE: We pass ``None`` for *iteration_stats* to the parent so that + # the upstream's ``_update_stats_from_output`` logs stats via its own + # ``RequestStateStats`` (req_state.stats) without touching the omni + # ``IterationStats`` object. Below we re-apply the same update using + # *native_stats* (the omni-specific ``RequestStateStats`` attached to + # ``OmniRequestState``), which correctly feeds the omni metrics path. + # If the upstream adds new logic in ``_update_stats_from_output`` (e.g. + # additional stat fields), this override must be kept in sync. super()._update_stats_from_output( req_state, engine_core_output, diff --git a/vllm_omni/entrypoints/cli/benchmark/main.py b/vllm_omni/entrypoints/cli/benchmark/main.py index d9d2268e8..391def31f 100644 --- a/vllm_omni/entrypoints/cli/benchmark/main.py +++ b/vllm_omni/entrypoints/cli/benchmark/main.py @@ -4,7 +4,7 @@ import argparse import typing from vllm.entrypoints.cli.types import CLISubcommand -from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG +from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG from vllm_omni.entrypoints.cli.benchmark.base import OmniBenchmarkSubcommandBase diff --git a/vllm_omni/entrypoints/cli/main.py b/vllm_omni/entrypoints/cli/main.py index 1a42b08f7..79129d732 100644 --- a/vllm_omni/entrypoints/cli/main.py +++ b/vllm_omni/entrypoints/cli/main.py @@ -23,7 +23,7 @@ def main(): if "VLLM_LOGGING_COLOR" not in os.environ: os.environ["VLLM_LOGGING_COLOR"] = "1" - from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup + from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG, cli_env_setup import vllm_omni.entrypoints.cli.benchmark.main import vllm_omni.entrypoints.cli.serve diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 723e1f236..5b6c7d9b3 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -17,7 +17,7 @@ from typing import Any import uvloop from vllm.entrypoints.cli.types import CLISubcommand from vllm.entrypoints.openai.cli_args import make_arg_parser, validate_parsed_serve_args -from vllm.entrypoints.utils import VLLM_SUBCMD_PARSER_EPILOG +from vllm.entrypoints.serve.utils.api_utils import VLLM_SUBCMD_PARSER_EPILOG from vllm.logger import init_logger from vllm_omni.entrypoints.cli.logo import log_logo diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index c11db925f..36c8201e0 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -34,7 +34,6 @@ from vllm.engine.protocol import EngineClient from vllm.entrypoints.anthropic.serving import AnthropicServingMessages from vllm.entrypoints.chat_utils import load_chat_template from vllm.entrypoints.launcher import serve_http, terminate_if_errored -from vllm.entrypoints.logger import RequestLogger from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer from vllm.entrypoints.openai.api_server import build_app as build_openai_app from vllm.entrypoints.openai.api_server import setup_server as setup_openai_server @@ -56,10 +55,7 @@ from vllm.entrypoints.openai.engine.protocol import ( ) from vllm.entrypoints.openai.models.protocol import BaseModelPath from vllm.entrypoints.openai.models.serving import OpenAIServingModels -from vllm.entrypoints.openai.orca_metrics import metrics_header from vllm.entrypoints.openai.responses.serving import OpenAIServingResponses -from vllm.entrypoints.openai.server_utils import get_uvicorn_log_config -from vllm.entrypoints.openai.utils import validate_json_request from vllm.entrypoints.pooling.classify.serving import ServingClassification from vllm.entrypoints.pooling.embed.serving import ServingEmbedding as OpenAIServingEmbedding from vllm.entrypoints.pooling.pooling.serving import ServingPooling @@ -71,6 +67,16 @@ from vllm.entrypoints.serve.disagg.serving import ServingTokens from vllm.entrypoints.serve.instrumentator.basic import base from vllm.entrypoints.serve.render.serving import OpenAIServingRender from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization +from vllm.entrypoints.serve.utils.api_utils import ( + load_aware_call, + process_lora_modules, + validate_json_request, + with_cancellation, +) +from vllm.entrypoints.serve.utils.error_response import create_error_response +from vllm.entrypoints.serve.utils.orca_metrics import metrics_header +from vllm.entrypoints.serve.utils.request_logger import RequestLogger +from vllm.entrypoints.serve.utils.server_utils import get_uvicorn_log_config from vllm.entrypoints.speech_to_text.realtime.serving import OpenAIServingRealtime from vllm.entrypoints.speech_to_text.transcription.serving import ( OpenAIServingTranscription, @@ -78,12 +84,6 @@ from vllm.entrypoints.speech_to_text.transcription.serving import ( from vllm.entrypoints.speech_to_text.translation.serving import ( OpenAIServingTranslation, ) -from vllm.entrypoints.utils import ( - create_error_response, - load_aware_call, - process_lora_modules, - with_cancellation, -) from vllm.logger import init_logger from vllm.tasks import POOLING_TASKS from vllm.tool_parsers import ToolParserManager diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 35cfafcff..2a7b7e269 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -69,8 +69,8 @@ from vllm.entrypoints.openai.parser.harmony_utils import ( parse_chat_output, ) 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.entrypoints.serve.utils.api_utils import should_include_usage +from vllm.entrypoints.serve.utils.tool_calls_utils import maybe_filter_parallel_tool_calls from vllm.inputs import PromptType from vllm.logger import init_logger from vllm.outputs import RequestOutput @@ -843,6 +843,27 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if value is not None and hasattr(obj, key): setattr(obj, key, value) + def _should_check_for_unstreamed_tool_arg_tokens( + self, + delta_message: Any, + output: Any, + ) -> bool: + """Check whether the streaming generator should flush unstreamed + tool-arg tokens at finish-time. + + This method was moved from OpenAIServingChat to the tool-parser layer + by upstream commit 9affc17a05. Omni's independently-maintained + ``chat_completion_stream_generator`` still calls it, so we keep a + local copy with the pre-9affc17a05 semantics. + """ + return ( + output.finish_reason is not None + and self.enable_auto_tools + and self.tool_parser is not None + and delta_message is not None + and delta_message.tool_calls + ) + def _build_sampling_params_list_from_request( self, request: ChatCompletionRequest, @@ -2794,6 +2815,16 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): if isinstance(ar_text, str) and ar_text.strip(): cot_output = ar_text + req_out = getattr(result, "request_output", None) + if req_out: + prompt_obj = getattr(req_out, "prompt", None) + if isinstance(prompt_obj, dict): + extra = prompt_obj.get("extra", {}) + if isinstance(extra, dict): + ar_text = extra.get("ar_generated_text") + if isinstance(ar_text, str) and ar_text.strip(): + cot_output = ar_text + return self._flatten_diffusion_images(images), stage_durations, peak_memory_mb, cot_output async def _stream_diffusion_image_chunks( diff --git a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py index bf6e9942f..b930578d6 100644 --- a/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py +++ b/vllm_omni/model_executor/models/minicpmo_4_5/minicpmo_4_5_omni_llm.py @@ -87,29 +87,17 @@ from vllm.config.multimodal import ( ImageDummyOptions, VideoDummyOptions, ) +from vllm.inputs import ModalityData, MultiModalDataDict from vllm.logger import init_logger from vllm.model_executor.models.interfaces import MultiModalEmbeddings, SupportsMRoPE, SupportsMultiModal, SupportsPP +from vllm.model_executor.models.utils import ( + _merge_multimodal_embeddings as merge_multimodal_embeddings, +) from vllm.model_executor.models.utils import ( flatten_bn, init_vllm_registered_model, maybe_prefix, ) - -try: - from vllm.model_executor.models.utils import merge_multimodal_embeddings -except ImportError: - from vllm.model_executor.models.utils import ( - _merge_multimodal_embeddings as merge_multimodal_embeddings, - ) -try: - from vllm.utils import flatten_2d_lists -except ImportError: - - def flatten_2d_lists(nested): - """vLLM >=0.18 removed vllm.utils.flatten_2d_lists; keep MiniCPM-o vision path working.""" - return [x for row in nested for x in row] - - from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.inputs import ( ImageItem, @@ -117,17 +105,6 @@ from vllm.multimodal.inputs import ( MultiModalKwargsItems, NestedTensors, ) - -# vllm >=0.21 moved ModalityData / MultiModalDataDict from -# vllm.multimodal.inputs to vllm.inputs (still re-exported from -# vllm.multimodal.parse). Fall back to the old location for older vllm. -try: - from vllm.inputs import ModalityData, MultiModalDataDict -except ImportError: - from vllm.multimodal.inputs import ( # type: ignore[no-redef] - ModalityData, - MultiModalDataDict, - ) from vllm.multimodal.parse import ( AudioItem, AudioProcessorItems, @@ -141,37 +118,19 @@ from vllm.multimodal.parse import ( VideoProcessorItems, ) from vllm.multimodal.processing import ( + BaseDummyInputsBuilder, BaseMultiModalProcessor, BaseProcessingInfo, PromptReplacement, PromptUpdate, PromptUpdateDetails, ) - -try: - from vllm.multimodal.processing import BaseDummyInputsBuilder -except ImportError: - from vllm.multimodal.profiling import BaseDummyInputsBuilder from vllm.sequence import IntermediateTensors def encode_tokens(tokenizer, prompt: str) -> list[int]: - """Tokenize ``prompt`` without adding special tokens. - - Prefer vllm's ``vllm.transformers_utils.tokenizer.encode_tokens`` when - available; falls back to the equivalent ``tokenizer.encode`` call so we - keep working across vllm versions that do not ship that helper. - - The import is performed lazily inside the function so the docs builder - (griffe / mkdocs strict mode) does not try to statically resolve the - upstream alias, which can fail across vllm releases that rename or move - the helper. - """ - try: - from vllm.transformers_utils.tokenizer import encode_tokens as _impl - except ImportError: - return tokenizer.encode(prompt, add_special_tokens=False) - return _impl(tokenizer, prompt) + """Tokenize ``prompt`` without adding special tokens.""" + return tokenizer.encode(prompt, add_special_tokens=False) from vllm.utils.tensor_schema import TensorSchema, TensorShape @@ -2855,7 +2814,7 @@ class MiniCPMO45OmniLLMProcessingInfo(BaseProcessingInfo): if self.ctx.tokenizer is not None: return self.ctx.tokenizer if not hasattr(self, "_lazy_tokenizer") or self._lazy_tokenizer is None: - from vllm.transformers_utils.tokenizer import cached_tokenizer_from_config + from vllm.tokenizers import cached_tokenizer_from_config self._lazy_tokenizer = cached_tokenizer_from_config(self.ctx.model_config) return self._lazy_tokenizer diff --git a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py index 11aeada9e..87b932886 100644 --- a/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py +++ b/vllm_omni/model_executor/models/moss_tts_nano/modeling_moss_tts_nano.py @@ -34,7 +34,10 @@ from vllm.config import VllmConfig from vllm.logger import init_logger from vllm_omni.model_executor.models.output_templates import OmniOutput -from vllm_omni.model_executor.models.utils import reinit_rotary_inv_freq +from vllm_omni.model_executor.models.utils import ( + reinit_rotary_inv_freq, + transformers_keys_to_ignore_compat, +) logger = init_logger(__name__) @@ -193,7 +196,10 @@ class MossTTSNanoForGeneration(nn.Module): logger.info("Loading MOSS-TTS-Nano LM from %s (dtype=%s)", self.model_path, tts_dtype) from transformers import AutoModelForCausalLM - with _hf_load_without_tp_warmup(): + # Two orthogonal trust_remote_code transformers fixes: + # - _hf_load_without_tp_warmup: tp_plan=None warmup crash (transformers 5.8.x) + # - transformers_keys_to_ignore_compat: keys list-vs-set (transformers 5.9) + with _hf_load_without_tp_warmup(), transformers_keys_to_ignore_compat(): lm = AutoModelForCausalLM.from_pretrained( self.model_path, trust_remote_code=True, @@ -233,7 +239,10 @@ class MossTTSNanoForGeneration(nn.Module): logger.info("Loading MOSS-Audio-Tokenizer-Nano from %s", codec_path) from transformers import AutoModel - with _hf_load_without_tp_warmup(): + # Two orthogonal trust_remote_code transformers fixes: + # - _hf_load_without_tp_warmup: tp_plan=None warmup crash (transformers 5.8.x) + # - transformers_keys_to_ignore_compat: keys list-vs-set (transformers 5.9) + with _hf_load_without_tp_warmup(), transformers_keys_to_ignore_compat(): audio_tokenizer = AutoModel.from_pretrained( codec_path, trust_remote_code=True, diff --git a/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py b/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py index 933236313..95ae62627 100644 --- a/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py +++ b/vllm_omni/model_executor/models/qwen3_omni/qwen3_moe.py @@ -5,7 +5,7 @@ import torch.nn.functional as F from torch import nn from vllm.config import VllmConfig from vllm.logger import init_logger -from vllm.model_executor.layers.fused_moe import FusedMoE +from vllm.model_executor.layers.fused_moe.layer import MoERunner from vllm.model_executor.layers.logits_processor import LogitsProcessor from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead from vllm.model_executor.models.qwen3_moe import ( @@ -149,14 +149,14 @@ class Qwen3MoeForCausalLM(_BaseQwen3MoeForCausalLM): # Set MoE hyperparameters for individual experts self.expert_weights = [] - self.moe_layers: list[FusedMoE] = [] + self.moe_layers: list[MoERunner] = [] example_layer = None for layer in self.model.layers: if isinstance(layer, PPMissingLayer): continue assert isinstance(layer, Qwen3MoeDecoderLayer) - if isinstance(layer.mlp, FusedMoE): + if isinstance(layer.mlp, MoERunner): example_layer = layer.mlp self.moe_layers.append(layer.mlp) diff --git a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py index b4f1d478c..b0294b725 100644 --- a/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py +++ b/vllm_omni/model_executor/models/qwen3_tts/prompt_embeds_builder.py @@ -903,12 +903,12 @@ class Qwen3TTSPromptEmbedsBuilder: * ``talker_prompt``: ``[prompt_len, hidden]`` prefill embedding. * ``trailing_text_hidden``: ``[T, hidden]`` queue of text-side - embeddings consumed one-per-decode-step (streaming mode) or a - single pad row (non-streaming). + embeddings consumed one-per-decode-step (streaming mode) + or a single pad row (non-streaming). * ``ref_code_len``: Number of ref codec frames consumed when the - Base task ran in in-context mode (``None`` otherwise). + Base task ran in in-context mode (``None`` otherwise). * ``ref_code``: The ``[T, Q]`` ref codec tensor used for ICL - (``None`` outside Base / ICL). + (``None`` outside Base / ICL). """ config = self._config talker_config = self._talker_config diff --git a/vllm_omni/model_executor/models/utils.py b/vllm_omni/model_executor/models/utils.py index 29d3be80f..5af5e00f1 100644 --- a/vllm_omni/model_executor/models/utils.py +++ b/vllm_omni/model_executor/models/utils.py @@ -1,9 +1,65 @@ from collections.abc import Callable +from contextlib import contextmanager import torch from vllm.model_executor.models.utils import maybe_prefix +@contextmanager +def transformers_keys_to_ignore_compat(): + """Make ``trust_remote_code`` weight loading robust to the transformers 5.9 + ``_keys_to_ignore_on_load_unexpected`` list-vs-set change. + + transformers 5.9 rewrote ``PreTrainedModel._adjust_missing_and_unexpected_keys`` + from ``(attr or []) + patterns`` (list concatenation) to + ``(attr or set()) | patterns`` (set union). Remote-code models such as + ``OpenMOSS-Team/MOSS-TTS-Nano`` still declare + ``_keys_to_ignore_on_load_unexpected`` as a *list*, so ``list | set`` raises + ``TypeError: unsupported operand type(s) for |: 'list' and 'set'`` and the + engine core dies during model load. + + Wrap any ``from_pretrained(..., trust_remote_code=True)`` call whose remote + code may declare the attribute as a list. The guard keeps such models + loadable regardless of which transformers version is installed, while + preserving the model's ignore patterns. + """ + try: + from transformers.modeling_utils import PreTrainedModel + except Exception: # pragma: no cover - transformers is always present here + yield + return + + orig = getattr(PreTrainedModel, "_adjust_missing_and_unexpected_keys", None) + if orig is None: + yield + return + + def _wrapper(self, *args, **kwargs): + try: + return orig(self, *args, **kwargs) + except TypeError as exc: + if "unsupported operand type" not in str(exc): + raise + attr = getattr(self, "_keys_to_ignore_on_load_unexpected", None) + if not isinstance(attr, (list, tuple)): + raise + # transformers >=5.9 combines patterns with ``set | ...``. The + # ignore-pattern combine runs before any mutation of ``loading_info`` + # (see PreTrainedModel._adjust_missing_and_unexpected_keys), so it is + # safe to coerce the list to a set and re-run the original once. + self._keys_to_ignore_on_load_unexpected = set(attr) + try: + return orig(self, *args, **kwargs) + finally: + self._keys_to_ignore_on_load_unexpected = attr + + PreTrainedModel._adjust_missing_and_unexpected_keys = _wrapper + try: + yield + finally: + PreTrainedModel._adjust_missing_and_unexpected_keys = orig + + def add_prefix_to_loaded_weights(weights: set[str], prefix: str) -> set[str]: """ Add a prefix to the names of the loaded weights. diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index c7cca20ea..d29af7314 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -6,6 +6,7 @@ and also outputs sampled tokens. from __future__ import annotations +import gc from collections.abc import Mapping from contextlib import nullcontext from copy import copy @@ -268,6 +269,63 @@ class GPUARModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin): self._capture_talker_mtp_graphs() return result + def shutdown(self) -> None: + """Release omni-specific GPU resources before upstream shutdown. + + Order of operations (must match upstream's expectation): + 1. Unfreeze Python GC so model weights are collected immediately + when self.model is set to None (upstream Worker.init_device + calls gc.freeze() / freeze_gc_heap()). + 2. Destroy omni-specific CUDA graphs (talker MTP) so references to + model parameters are released before self.model = None. + 3. Clear GPU-side buffers (input_ids, inputs_embeds) and per-request + caches that may hold GPU tensor references. + 4. Call CUDAGraphWrapper.clear_all_graphs() unconditionally (not just + on ROCm) to ensure all CUDA graphs including talker MTP are + released before model weight teardown. + 5. Call BreakableCUDAGraphWrapper.clear_all_graphs() as well, to + match the upstream ROCm-only pattern but also protect CUDA. + 6. Delegate to upstream GPUModelRunner.shutdown() which sets + self.model = None, clears KV caches, resets workspace, etc. + + This prevents abrupt GPU memory release during EngineCore subprocess + exit that can trigger GPU OOM signals when the parent process + concurrently cleans up its own GPU state. + """ + # 1. Unfreeze GC so model weights and GPU tensors are collected + # immediately when references are dropped (upstream Worker.shutdown + # also does this before any teardown). + gc.unfreeze() + + # 2. Destroy talker MTP CUDA graph wrapper to release captured graphs. + if hasattr(self, "talker_mtp") and self.talker_mtp is not None: + self.talker_mtp = None + self.has_talker_mtp = False + + # 3. Clear GPU-side buffers (small tensors, but every MiB helps). + if hasattr(self, "input_ids") and self.input_ids is not None: + self.input_ids = None + if hasattr(self, "inputs_embeds") and self.inputs_embeds is not None: + self.inputs_embeds = None + + # 4. Clear per-request caches that may hold GPU tensor references. + if hasattr(self, "_downstream_payload_cache"): + self._downstream_payload_cache.clear() + if hasattr(self, "model_intermediate_buffer"): + self.model_intermediate_buffer.clear() + + # 5. Release all CUDA graphs unconditionally (upstream only does this + # on ROCm; on CUDA the graphs are only freed by Python GC during + # interpreter shutdown, which is too late to prevent memory spikes). + from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphWrapper + from vllm.compilation.cuda_graph import CUDAGraphWrapper + + CUDAGraphWrapper.clear_all_graphs() + BreakableCUDAGraphWrapper.clear_all_graphs() + + # 6. Delegate to upstream shutdown (model = None, KV caches, workspace). + super().shutdown() + def _capture_talker_mtp_graphs(self) -> None: from vllm.compilation.cuda_graph import CUDAGraphWrapper diff --git a/vllm_omni/worker/gpu_ar_worker.py b/vllm_omni/worker/gpu_ar_worker.py index 50ad70c57..08551ad69 100644 --- a/vllm_omni/worker/gpu_ar_worker.py +++ b/vllm_omni/worker/gpu_ar_worker.py @@ -9,12 +9,12 @@ 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 init_worker_distributed_environment -from vllm.v1.worker.utils import request_memory from vllm.v1.worker.workspace import init_workspace_manager from vllm_omni.diffusion.data import OmniACK, OmniSleepTask, OmniWakeTask from vllm_omni.worker.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_ar_model_runner import GPUARModelRunner +from vllm_omni.worker.memory_utils import request_memory_tolerant from vllm_omni.worker.mixins import OmniWorkerMixin logger = init_logger(__name__) @@ -84,7 +84,7 @@ class GPUARWorker(OmniWorkerMixin, OmniGPUWorkerBase): # take current memory snapshot self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) - self.requested_memory = request_memory(init_snapshot, self.cache_config) + self.requested_memory = request_memory_tolerant(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: diff --git a/vllm_omni/worker/gpu_generation_model_runner.py b/vllm_omni/worker/gpu_generation_model_runner.py index 793228158..bc3d6ea2d 100644 --- a/vllm_omni/worker/gpu_generation_model_runner.py +++ b/vllm_omni/worker/gpu_generation_model_runner.py @@ -718,6 +718,7 @@ class GPUGenerationModelRunner(OmniGPUModelRunner, OmniConnectorModelRunnerMixin cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1]) self.query_start_loc.copy_to_gpu() self.input_batch.block_table.commit_block_table(num_reqs_padded) diff --git a/vllm_omni/worker/gpu_generation_worker.py b/vllm_omni/worker/gpu_generation_worker.py index a356f03ad..2a85436a0 100644 --- a/vllm_omni/worker/gpu_generation_worker.py +++ b/vllm_omni/worker/gpu_generation_worker.py @@ -9,11 +9,11 @@ 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 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.base import OmniGPUWorkerBase from vllm_omni.worker.gpu_generation_model_runner import GPUGenerationModelRunner +from vllm_omni.worker.memory_utils import request_memory_tolerant from vllm_omni.worker.mixins import OmniWorkerMixin logger = init_logger(__name__) @@ -83,7 +83,7 @@ class GPUGenerationWorker(OmniWorkerMixin, OmniGPUWorkerBase): # take current memory snapshot self.init_snapshot = init_snapshot = MemorySnapshot(device=self.device) - self.requested_memory = request_memory(init_snapshot, self.cache_config) + self.requested_memory = request_memory_tolerant(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: diff --git a/vllm_omni/worker/gpu_model_runner.py b/vllm_omni/worker/gpu_model_runner.py index 120c7a565..349e1da04 100644 --- a/vllm_omni/worker/gpu_model_runner.py +++ b/vllm_omni/worker/gpu_model_runner.py @@ -1006,6 +1006,7 @@ class OmniGPUModelRunner(GPUModelRunner): cum_num_tokens = self._get_cumsum_and_arange(num_scheduled_tokens, self.query_pos.np) self.query_start_loc.np[1 : num_reqs + 1] = cum_num_tokens + self.query_start_loc.np[num_reqs + 1 : num_reqs_padded + 1].fill(cum_num_tokens[-1]) self.query_start_loc.copy_to_gpu() self.input_batch.block_table.commit_block_table(num_reqs_padded) diff --git a/vllm_omni/worker/memory_utils.py b/vllm_omni/worker/memory_utils.py new file mode 100644 index 000000000..321114f88 --- /dev/null +++ b/vllm_omni/worker/memory_utils.py @@ -0,0 +1,55 @@ +"""GPU memory utilities for vLLM Omni workers. + +Includes a tolerant version of the upstream request_memory() that handles +multi-stage GPU sharing by capping the memory budget to available free +memory instead of raising ValueError. +""" + +from __future__ import annotations + +import math + +from vllm.config import CacheConfig +from vllm.logger import init_logger +from vllm.utils.mem_utils import MemorySnapshot, format_gib + +logger = init_logger(__name__) + + +def request_memory_tolerant( + init_snapshot: MemorySnapshot, + cache_config: CacheConfig, +) -> int: + """Calculate the amount of memory required for this stage. + + Like upstream ``request_memory()`` but tolerates multi-stage GPU sharing: + if ``free_memory < requested_memory`` (because another stage on the same + GPU has already consumed memory), caps the requested budget to the actual + free memory instead of raising ``ValueError``. The downstream + ``OmniGPUWorkerBase.determine_available_memory()`` already does per-process + NVML accounting and correctly computes the KV cache budget regardless. + + Logs a warning when the budget is capped so operators can detect + under-provisioned GPU memory. + """ + requested_memory = math.ceil(init_snapshot.total_memory * cache_config.gpu_memory_utilization) + + if init_snapshot.free_memory < requested_memory: + capped = init_snapshot.free_memory + logger.warning( + "Free memory on device %s (%s/%s GiB) on startup is less than " + "desired GPU memory utilization (%.2f, %s GiB). " + "Capping requested memory to available free memory (%s GiB). " + "This is expected when multiple Omni stages share a GPU; " + "the per-process NVML accounting in determine_available_memory() " + "will compute the correct KV cache budget.", + init_snapshot.device_, + format_gib(init_snapshot.free_memory), + format_gib(init_snapshot.total_memory), + cache_config.gpu_memory_utilization, + format_gib(requested_memory), + format_gib(capped), + ) + return capped + + return requested_memory