From d3c47efc4f818931f29e5f1567c94c454561631c Mon Sep 17 00:00:00 2001 From: Zhou Taichang Date: Sun, 12 Jul 2026 23:39:53 +0800 Subject: [PATCH] [Rebase] Rebase to vllm v0.25.0 (#5042) Signed-off-by: tzhouam Co-authored-by: Claude Opus 4.8 --- docker/Dockerfile.ci | 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 +- tests/ar_diffusion/test_kv_cache.py | 12 +--- .../test_async_omni_engine_stage_init.py | 51 ++++++++++++++ .../test_qwen3_omni_realtime_websocket.py | 20 +++++- .../openai_api/test_serving_chat_speaker.py | 4 +- .../entrypoints/test_stream_finish_reason.py | 2 +- .../voxcpm2/test_talker_state_eviction.py | 24 +++++-- vllm_omni/core/sched/omni_ar_scheduler.py | 21 ++++-- .../core/sched/omni_generation_scheduler.py | 22 +++--- .../model_loader/diffusers_loader.py | 65 +++++++++++++++++- .../gr00t/modeling/processing_gr00t_n1d7.py | 8 ++- vllm_omni/engine/async_omni_engine.py | 12 ++-- vllm_omni/engine/output_processor.py | 18 +++++ vllm_omni/engine/stage_engine_startup.py | 58 +++++++++++++++- vllm_omni/entrypoints/openai/api_server.py | 46 ++++++------- .../openai/serving_audio_generate.py | 2 +- vllm_omni/entrypoints/openai/serving_chat.py | 21 +++++- .../entrypoints/openai/serving_speech.py | 2 +- .../ar_diffusion/kv_cache/paged.py | 17 +---- vllm_omni/patch.py | 67 +++++++++++++++++++ vllm_omni/worker/gpu_ar_model_runner.py | 12 ++++ 25 files changed, 415 insertions(+), 109 deletions(-) diff --git a/docker/Dockerfile.ci b/docker/Dockerfile.ci index 8b19ed678..f4444f8ac 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.24.0 +ARG VLLM_BASE_TAG=v0.25.0 FROM ${VLLM_BASE_IMAGE}:${VLLM_BASE_TAG} ARG APP_DIR=/workspace/vllm-omni WORKDIR ${APP_DIR} diff --git a/docs/getting_started/installation/gpu.md b/docs/getting_started/installation/gpu.md index 4ef35398c..e21a11cee 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, 0.22.0, 0.23.0, and 0.24.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, 0.23.0, 0.24.0, and 0.25.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 83cbcd68d..d4808139e 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.24.0 --torch-backend=auto +uv pip install vllm==0.25.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.24.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library ```bash -uv pip install vllm==0.24.0 --torch-backend=auto +uv pip install vllm==0.25.0 --torch-backend=auto ``` -The 0.24.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. +The 0.25.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.24.0 +git checkout v0.25.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 13.0 (the default for v0.24.0; the wheel filename has no `+cu130` suffix) -export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.24.0/vllm-0.24.0-cp38-abi3-manylinux_2_28_x86_64.whl +#For CUDA 13.0 (the default for v0.25.0; the wheel filename has no `+cu130` suffix) +export VLLM_PRECOMPILED_WHEEL_LOCATION=https://github.com/vllm-project/vllm/releases/download/v0.25.0/vllm-0.25.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.24.0 \ + vllm/vllm-omni:v0.25.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 af88e1e0c..4a64e8723 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.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722 +uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723 ``` #### 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.24.0 release version of the library +If you do not need to modify source code of vLLM, you can directly install the stable 0.25.0 release version of the library ```bash -uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722 +uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723 ``` -The pre-built 0.24.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.25.0 vLLM wheel targets ROCm 7.2.3. 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.24.0 +git checkout v0.25.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.24.0 \ + vllm/vllm-omni-rocm:v0.25.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.24.0 + vllm/vllm-omni-rocm:v0.25.0 ``` # --8<-- [end:pre-built-images] diff --git a/docs/getting_started/quickstart.md b/docs/getting_started/quickstart.md index 2c7f278ca..914f1f7b4 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.24.0 --torch-backend=auto +uv pip install vllm==0.25.0 --torch-backend=auto # On ROCm -uv pip install vllm==0.24.0+rocm722 --extra-index-url https://wheels.vllm.ai/rocm/0.24.0/rocm722 +uv pip install vllm==0.25.0+rocm723 --extra-index-url https://wheels.vllm.ai/rocm/0.25.0/rocm723 git clone https://github.com/vllm-project/vllm-omni.git cd vllm-omni @@ -34,7 +34,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.24.0` and vLLM Omni `0.24.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.25.0` and vLLM Omni `0.25.0`, as vLLM Omni no longer hijacks the vLLM entrypoint. Updating vLLM should resolve this issue. ## Offline Inference diff --git a/tests/ar_diffusion/test_kv_cache.py b/tests/ar_diffusion/test_kv_cache.py index cbb9656e5..beb575a0b 100644 --- a/tests/ar_diffusion/test_kv_cache.py +++ b/tests/ar_diffusion/test_kv_cache.py @@ -9,14 +9,9 @@ KV stack on CPU (block bookkeeping only, no GPU tensors). import pytest import torch from vllm.v1.kv_cache_interface import KVCacheSpecKind, get_kv_cache_spec_kind +from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry from vllm.v1.request import RequestStatus -try: - from vllm.v1.kv_cache_spec_registry import KVCacheSpecRegistry -except ModuleNotFoundError: - KVCacheSpecRegistry = None - from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map - from vllm_omni.experimental.ar_diffusion.kv_cache import ( ARDiffusionKVCache, ARDiffusionKVConfig, @@ -53,10 +48,7 @@ def test_spec_registration_resolves_to_chunk_window_manager(): # Without explicit registration the MRO walk would fall back to the parent # SlidingWindowManager; assert the subclass manager wins. spec = make_spec() - if KVCacheSpecRegistry is not None: - assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager - else: - assert spec_manager_map[type(spec)] is ChunkWindowManager + assert KVCacheSpecRegistry.get_manager_class(spec) is ChunkWindowManager def test_spec_kind_is_sliding_window(): diff --git a/tests/engine/test_async_omni_engine_stage_init.py b/tests/engine/test_async_omni_engine_stage_init.py index ff467dc65..d5a50cf29 100644 --- a/tests/engine/test_async_omni_engine_stage_init.py +++ b/tests/engine/test_async_omni_engine_stage_init.py @@ -874,3 +874,54 @@ def test_build_engine_args_dict_uses_diffusion_attention_config_key(): assert "attention_config" not in engine_args_dict assert engine_args_dict["diffusion_attention_config"].default.backend == "FLASH_ATTN" + + +def test_omni_master_server_allocates_globally_unique_route_ports(monkeypatch): + """Regression: two stages must never draw the same ZMQ port. + + ``get_open_ports_list`` only dedups within a single call, so per-route + allocation used to let a later stage reuse an earlier stage's port. The + second engine to ``bind()`` then died with ``zmq.error.ZMQError: Address + already in use`` (flaky multi-stage startup, e.g. Qwen3-Omni thinker/talker/ + code2wav). ``OmniMasterServer`` now dedups every port it hands out. + """ + import itertools + + from vllm_omni.engine import stage_engine_startup as ses + + # A colliding prefix (repeats + the master port 9000) followed by an endless + # fresh stream, so the only way to succeed is to redraw the collisions. + supply = itertools.chain([9000, 9000, 9001, 9001, 9000], itertools.count(9002)) + + def fake_get_open_ports_list(count): + return [next(supply) for _ in range(count)] + + monkeypatch.setattr(ses, "get_open_ports_list", fake_get_open_ports_list) + + server = ses.OmniMasterServer( + master_address="127.0.0.1", + master_port=9000, # seed: a route must not reuse the registration port + stage_ids=[0, 1, 2], + ) + + ports = [] + for sid in (0, 1, 2): + alloc = server.get_allocation(sid) + for addr in ( + alloc.handshake_bind_address, + alloc.input_bind_address, + alloc.output_bind_address, + ): + ports.append(ses._port_from_zmq_address(addr)) + + assert len(ports) == len(set(ports)), f"duplicate route ports allocated: {ports}" + assert 9000 not in ports, "route reused the master registration port" + + +def test_port_from_zmq_address_parsing(): + from vllm_omni.engine.stage_engine_startup import _port_from_zmq_address + + assert _port_from_zmq_address("tcp://127.0.0.1:34277") == 34277 + assert _port_from_zmq_address(None) is None + assert _port_from_zmq_address("ipc:///tmp/sock") is None + assert _port_from_zmq_address("tcp://host:not-a-port") is None diff --git a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py index 15e257959..50f8eef91 100644 --- a/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py +++ b/tests/entrypoints/openai_api/test_qwen3_omni_realtime_websocket.py @@ -198,12 +198,28 @@ def _assert_realtime_smoke(result: dict) -> None: assert result["output_sample_rate"] > 0 -def _assert_realtime_accuracy(result: dict) -> None: +def _assert_realtime_accuracy( + result: dict, + whisper_model_size: str = "large-v3", +) -> None: + """Assert that whisper transcription of audio output matches model text. + + Args: + result: Roundtrip result dict from ``_run_realtime_audio_roundtrip``. + whisper_model_size: Whisper model used to transcribe the generated audio + for the accuracy check. Defaults to ``large-v3``: the default + ``small`` model mishears short Chinese TTS clips (observed: + 北京→韦京 and a dropped leading sentence, sim=0.443), which + caused spurious sim<0.8 failures under async_chunk codec + variability even though audio generation was correct. large-v3 + transcribes these clips reliably, so a failure here now points + at the model, not the ASR grader. + """ final_text = (result["transcription_text"] or "").strip() assert final_text, "Expected non-empty transcription (model text stream)" wav_out = _wav_bytes_from_pcm16(result["output_pcm"], result["output_sample_rate"]) - whisper_text = convert_audio_bytes_to_text(wav_out).strip() + whisper_text = convert_audio_bytes_to_text(wav_out, model_size=whisper_model_size).strip() assert whisper_text, "Whisper returned empty string for synthesized output audio" sim = cosine_similarity_text(whisper_text.lower(), final_text.lower()) diff --git a/tests/entrypoints/openai_api/test_serving_chat_speaker.py b/tests/entrypoints/openai_api/test_serving_chat_speaker.py index 27839ee23..b0977c0c2 100644 --- a/tests/entrypoints/openai_api/test_serving_chat_speaker.py +++ b/tests/entrypoints/openai_api/test_serving_chat_speaker.py @@ -87,8 +87,8 @@ def test_create_chat_completion_converts_value_error_to_error_response(mocker: M serving_chat.chat_template = None serving_chat.chat_template_content_format = "string" serving_chat.default_chat_template_kwargs = {} - serving_chat.openai_serving_render = mocker.MagicMock() - serving_chat.openai_serving_render.validate_chat_template.return_value = None + serving_chat.online_renderer = mocker.MagicMock() + serving_chat.online_renderer.validate_chat_template.return_value = None serving_chat._effective_chat_template_kwargs = mocker.MagicMock(return_value={}) serving_chat._preprocess_chat = mocker.AsyncMock( side_effect=ValueError("Invalid speaker 'uncle_fu'. Supported: ethan, vivian") diff --git a/tests/entrypoints/test_stream_finish_reason.py b/tests/entrypoints/test_stream_finish_reason.py index 52cc32215..ebbca1917 100644 --- a/tests/entrypoints/test_stream_finish_reason.py +++ b/tests/entrypoints/test_stream_finish_reason.py @@ -144,7 +144,7 @@ def _build_serving_chat(): engine_client=mock_engine, models=models, response_role="assistant", - openai_serving_render=mock_render, + online_renderer=mock_render, request_logger=None, chat_template=None, chat_template_content_format="auto", diff --git a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py index 28da0b1bc..f5a94376a 100644 --- a/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py +++ b/tests/model_executor/models/voxcpm2/test_talker_state_eviction.py @@ -52,6 +52,10 @@ def _make_bare_talker(): talker._cfm_buffers = None talker._last_audio_output_req_ids = [] talker._batched_fsq_fusion_max_batch = 32 + # Added in VoxCPM2TalkerForConditionalGeneration.__init__; set to + # default (False) to avoid AttributeError when test exercises forward + # path that checks this flag. + talker._enable_unified_decode_graph = False return talker @@ -329,23 +333,27 @@ class TestDecodeBatchContract: for (expected_res, expected_meta), (actual_res, actual_meta) in zip(expected, actual, strict=True): torch.testing.assert_close(actual_res, expected_res) - torch.testing.assert_close(actual_meta["new_lm_hidden"], expected_meta["new_lm_hidden"]) + torch.testing.assert_close(actual_meta.new_lm_hidden, expected_meta.new_lm_hidden) assert [state.decode_step_count for state in batch_states] == [1, 1] def test_batch_decode_matches_per_request_state_updates(self) -> None: + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + _DecodeResidualMeta, + ) + seq_talker = self._make_decode_talker() batch_talker = self._make_decode_talker() seq_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)] batch_states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)] metas = [ - {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])}, - {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])}, + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])), + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])), ] batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]]) for state, meta, res_out in zip(seq_states, metas, batch_out, strict=True): - seq_talker._finish_decode(state, meta, res_out.unsqueeze(0), torch.device("cpu")) + seq_talker._finish_decode(state, meta, res_out.unsqueeze(0)) batch_talker._finish_decode_batch( [(state, False, meta) for state, meta in zip(batch_states, metas, strict=True)], @@ -360,12 +368,16 @@ class TestDecodeBatchContract: torch.testing.assert_close(actual.last_audio_patch_gpu, expected.last_audio_patch_gpu) def test_batch_decode_preserves_per_request_cfm_rng_order(self) -> None: + from vllm_omni.model_executor.models.voxcpm2.voxcpm2_talker import ( + _DecodeResidualMeta, + ) + talker = self._make_decode_talker() talker._enable_batched_cfm = False states = [self._make_decode_state("req-0", 0.0), self._make_decode_state("req-1", 10.0)] metas = [ - {"new_lm_hidden": torch.tensor([[1.0, 2.0, 3.0, 4.0]])}, - {"new_lm_hidden": torch.tensor([[5.0, 6.0, 7.0, 8.0]])}, + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[1.0, 2.0, 3.0, 4.0]])), + _DecodeResidualMeta(new_lm_hidden=torch.tensor([[5.0, 6.0, 7.0, 8.0]])), ] batch_out = torch.tensor([[0.5, 1.0, 1.5, 2.0], [2.5, 3.0, 3.5, 4.0]]) diff --git a/vllm_omni/core/sched/omni_ar_scheduler.py b/vllm_omni/core/sched/omni_ar_scheduler.py index 217c81543..74d09ab60 100644 --- a/vllm_omni/core/sched/omni_ar_scheduler.py +++ b/vllm_omni/core/sched/omni_ar_scheduler.py @@ -314,13 +314,6 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None - kv_connector_stats: KVConnectorStats | None = ( - kv_connector_output.kv_connector_stats if kv_connector_output else None - ) - if kv_connector_stats and self.connector: - kv_stats = self.connector.get_kv_connector_stats() - if kv_stats: - kv_connector_stats = kv_connector_stats.aggregate(kv_stats) failed_kv_load_req_ids = None if kv_connector_output and kv_connector_output.invalid_block_ids: @@ -563,6 +556,20 @@ class OmniARScheduler(OmniSchedulerMixin, VLLMScheduler): if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) + # Worker-side KV connector stats from the model runner output. + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) + if self.connector: + # Scheduler-side KV connector stats collected after connector update. + scheduler_kv_connector_stats = self.connector.get_kv_connector_stats() + if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty(): + kv_connector_stats = ( + kv_connector_stats.aggregate(scheduler_kv_connector_stats) + if kv_connector_stats is not None + else scheduler_kv_connector_stats + ) + # collect KV cache events from KV cache manager events = self.kv_cache_manager.take_events() diff --git a/vllm_omni/core/sched/omni_generation_scheduler.py b/vllm_omni/core/sched/omni_generation_scheduler.py index 42caea252..973b8c55a 100644 --- a/vllm_omni/core/sched/omni_generation_scheduler.py +++ b/vllm_omni/core/sched/omni_generation_scheduler.py @@ -431,14 +431,6 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): outputs: dict[int, list[EngineCoreOutput]] = defaultdict(list) spec_decoding_stats: SpecDecodingStats | None = None - kv_connector_stats: KVConnectorStats | None = ( - kv_connector_output.kv_connector_stats if kv_connector_output else None - ) - # Merge connector-side stats (align with v0.14.0) - if kv_connector_stats and self.connector: - kv_stats = self.connector.get_kv_connector_stats() - if kv_stats: - kv_connector_stats = kv_connector_stats.aggregate(kv_stats) failed_kv_load_req_ids = None if kv_connector_output and getattr(kv_connector_output, "invalid_block_ids", None): @@ -646,6 +638,20 @@ class OmniGenerationScheduler(OmniSchedulerMixin, VLLMScheduler): if kv_connector_output: self._update_from_kv_xfer_finished(kv_connector_output) + # Worker-side KV connector stats from the model runner output. + kv_connector_stats: KVConnectorStats | None = ( + kv_connector_output.kv_connector_stats if kv_connector_output else None + ) + if self.connector: + # Scheduler-side KV connector stats collected after connector update. + scheduler_kv_connector_stats = self.connector.get_kv_connector_stats() + if scheduler_kv_connector_stats is not None and not scheduler_kv_connector_stats.is_empty(): + kv_connector_stats = ( + kv_connector_stats.aggregate(scheduler_kv_connector_stats) + if kv_connector_stats is not None + else scheduler_kv_connector_stats + ) + # Collect and publish KV cache events (align with v0.14.0) events = self.kv_cache_manager.take_events() if self.connector is not None: diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index 943c55a52..d8b9c90c5 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -392,16 +392,77 @@ class DiffusersPipelineLoader: return model.eval() + @staticmethod + def _has_online_quant(model: nn.Module) -> bool: + """Whether any layer uses an online-quant method that defers weight + materialization onto the ``meta`` device (upstream vLLM + ``uses_meta_device=True``, e.g. online FP8).""" + for module in model.modules(): + quant_method = getattr(module, "quant_method", None) + if getattr(quant_method, "uses_meta_device", False): + return True + return False + def _process_weights_after_loading(self, model: nn.Module, target_device: torch.device) -> None: """Process weights after loading for quantization methods. This handles vLLM's quantization methods that need to process weights after loading (e.g., FP8 online quantization from BF16/FP16 weights). """ + # Newer upstream vLLM online-quant methods (uses_meta_device=True) create + # weights on the ``meta`` device and materialize them just-in-time as each + # layer's weights finish loading (via the layerwise online-process loader). + # Any "straggler" layers whose weights were not fully materialized during + # load (padded / partially-loaded layers) remain on ``meta``. Upstream's + # base_loader calls finalize_layerwise_processing() to materialize them; + # the diffusion loader must mirror that, otherwise the module.to() below + # raises "Cannot copy out of meta tensor; no data!". This whole meta-device + # handling is gated on online quant actually being in use, so that the + # proven code path for everything else (in particular FSDP/HSDP-sharded + # params, whose per-parameter .data cannot be cross-device reassigned) is + # left untouched. Import lazily so older vLLM (no meta-device quant) is + # unaffected. + has_online_quant = self._has_online_quant(model) + if has_online_quant: + from vllm.model_executor.model_loader.reload.layerwise import ( + finalize_layerwise_processing, + ) + + # model_config is only dereferenced by finalize for vLLM Attention / + # MLAAttention layers; diffusion DiT models use their own attention and + # have none, so passing None is safe here. + finalize_layerwise_processing(model, model_config=None) + for _, module in model.named_modules(): quant_method = getattr(module, "quant_method", None) - if quant_method is not None and isinstance(quant_method, QuantizeMethodBase): - # Move module to target device for processing if needed + if quant_method is None or not isinstance(quant_method, QuantizeMethodBase): + continue + + if has_online_quant: + # Online quant may leave straggler params on the ``meta`` device. + # Move only real (non-meta) params onto the target device for + # processing and restore them afterward, mirroring upstream vLLM's + # device_loading_context — a blanket module.to(target_device) would + # raise NotImplementedError on meta params. Online quant initializes + # on the accelerator, so params are normally already on the target + # device and this loop is a no-op move; the point is to skip meta. + original_devices: dict[str, torch.device] = {} + for name, param in module.named_parameters(): + if param.device.type != "meta" and param.device != target_device: + original_devices[name] = param.device + param.data = param.data.to(target_device) + + quant_method.process_weights_after_loading(module) + + # Restore pre-existing params to their original device; leave any + # newly created (e.g. quantized) params on the target device. + for name, param in module.named_parameters(): + if name in original_devices: + param.data = param.data.to(original_devices[name]) + else: + # No meta params possible here. Preserve the original FSDP/HSDP-aware + # whole-module move (module.to()), which correctly handles sharded + # DTensor params that per-parameter .data reassignment cannot. module_device = next(module.parameters(), None) if module_device is not None: module_device = module_device.device diff --git a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py index ba128f450..ade5eac83 100644 --- a/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py +++ b/vllm_omni/diffusion/models/gr00t/modeling/processing_gr00t_n1d7.py @@ -19,6 +19,7 @@ from transformers.utils import cached_file from vllm.logger import init_logger from vllm_omni.diffusion.models.gr00t.configs.embodiment_configs import ModalityConfig +from vllm_omni.diffusion.models.gr00t.configs.gr00t_n1d7 import Gr00tN1d7Config from vllm_omni.diffusion.models.gr00t.dataio.embodiment_tags import EmbodimentTag from vllm_omni.diffusion.models.gr00t.dataio.state_action.state_action_processor import StateActionProcessor from vllm_omni.diffusion.models.gr00t.dataio.utils import parse_modality_configs, to_json_serializable @@ -552,4 +553,9 @@ class Gr00tN1d7Processor(ProcessorMixin): return cls(**processor_kwargs, transformers_loading_kwargs=transformers_loading_kwargs) -AutoProcessor.register("Gr00tN1d7", Gr00tN1d7Processor) +# transformers >= 5.5 requires the config *class* (not a model_type string) as the +# first arg to AutoProcessor.register — it does `key.__module__` internally. Passing +# the string "Gr00tN1d7" raises AttributeError: 'str' object has no attribute +# '__module__'. Match the AutoModel.register(Gr00tN1d7Config, ...) call in +# modeling/gr00t_n1d7.py. +AutoProcessor.register(Gr00tN1d7Config, Gr00tN1d7Processor) diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index ec9a9698c..ce9662884 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1568,16 +1568,18 @@ class AsyncOmniEngine: # 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. + # CuMemAllocator.sleep() is NOT idempotent — calling it on already- + # slept entries causes CUDA_ERROR_INVALID_VALUE at cumem_allocator + # cuMemRelease (double-free of the memory handle). Use release_pools() + # instead, which is the designed cleanup path: it drops MemPool refs + # and lets the destructor/free path handle asleep entries correctly + # (returns a null handle so the C extension skips unmap/release). 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() + allocator.release_pools() logger.debug("[AsyncOmniEngine] Released CuMem memory pool during shutdown") except Exception: pass diff --git a/vllm_omni/engine/output_processor.py b/vllm_omni/engine/output_processor.py index c9044e29e..71303269f 100644 --- a/vllm_omni/engine/output_processor.py +++ b/vllm_omni/engine/output_processor.py @@ -28,6 +28,12 @@ from vllm_omni.outputs import OmniRequestOutput logger = init_logger(__name__) +# Keys whose values may arrive as scalar tensors (e.g. sample rate) but +# represent metadata, not concatenatable data. When these appear in +# ``MultimodalPayload.tensors`` they must be relocated to ``.metadata`` +# before consolidation so the tensor concatenation path does not attempt +# a dimensioned cat on a 0-d tensor. +_METADATA_TENSOR_KEYS: frozenset[str] = frozenset({"sr", "sample_rate", "audio_sample_rate"}) # --------------------------------------------------------------------------- # Module-level helpers @@ -212,6 +218,18 @@ class OmniRequestState(RequestState): modality = OutputModality.TEXT strategy = get_accumulation_strategy(modality) + # Move known metadata scalar tensors (e.g. sample rate) from + # tensors to metadata so they are handled by the REPLACE path + # below instead of failing tensor concatenation. + for key in _METADATA_TENSOR_KEYS: + v = self.mm_accumulated.tensors.pop(key, None) + if v is not None: + if isinstance(v, list): + # Keep the last value (REPLACE semantics) + self.mm_accumulated.metadata[key] = v[-1] + else: + self.mm_accumulated.metadata[key] = v + try: for k, v in list(self.mm_accumulated.tensors.items()): if isinstance(v, list) and v and isinstance(v[0], torch.Tensor): diff --git a/vllm_omni/engine/stage_engine_startup.py b/vllm_omni/engine/stage_engine_startup.py index 8f6c199ec..7a55bb93f 100644 --- a/vllm_omni/engine/stage_engine_startup.py +++ b/vllm_omni/engine/stage_engine_startup.py @@ -138,6 +138,21 @@ class StageReplicaResources: # --------------------------------------------------------------------------- +def _port_from_zmq_address(address: str | None) -> int | None: + """Extract the TCP port from a ``tcp://host:port`` ZMQ address, else ``None``. + + Non-TCP transports (``ipc://``, ``inproc://``) and unparsable values return + ``None`` so callers can skip seeding them into the port-dedup set. + """ + if not address: + return None + tail = address.rsplit(":", 1)[-1] + try: + return int(tail) + except ValueError: + return None + + class OmniMasterServer: """Registration server for single-stage engine startup.""" @@ -158,10 +173,25 @@ class OmniMasterServer: self._stage_configs: dict[StageRoute, Any] = {} self._stage_coordinator_addresses: dict[StageRoute, StageCoordinatorAddresses] = {} self._stage_config_events: dict[StageRoute, threading.Event] = {} + # Ports already handed out by *this* server. ``get_open_ports_list`` only + # guarantees uniqueness *within* a single call; a per-route call cannot + # see ports drawn for other routes, so two stages/replicas can draw the + # same ephemeral port and the second engine to ``bind()`` it dies with + # ``zmq.error.ZMQError: Address already in use``. Multi-stage models + # (e.g. Qwen3-Omni: thinker/talker/code2wav) allocate many routes and hit + # this intermittently. We dedup every allocated port against this set, + # seeded with the registration port (and coordinator ROUTER port) that + # are already bound on the same host. + self._allocated_ports: set[int] = set() + if master_port: + self._allocated_ports.add(int(master_port)) # Coordinator ROUTER address echoed back in every registration reply # so OmniCoordClientForStage knows where to connect from inside the # engine subprocess. self._coordinator_router_address = coordinator_router_address + coord_port = _port_from_zmq_address(coordinator_router_address) + if coord_port is not None: + self._allocated_ports.add(coord_port) # Fires only for *newly assigned* (auto-assigned) replicas, not for # head-side pre-allocated slots that already have head-side clients. self._on_register = on_register @@ -235,7 +265,7 @@ class OmniMasterServer: self._stage_config_events[route] = threading.Event() self._stage_coordinator_addresses[route] = StageCoordinatorAddresses() - hs_port, inp_port, out_port = get_open_ports_list(count=3) + hs_port, inp_port, out_port = self._alloc_unique_ports(3) alloc = StageAllocation( handshake_bind_address=f"tcp://{self._address}:{hs_port}", handshake_connect_address=f"tcp://{self._address}:{hs_port}", @@ -247,6 +277,32 @@ class OmniMasterServer: self._stage_routes[route] = alloc return alloc + def _alloc_unique_ports(self, count: int) -> list[int]: + """Return ``count`` open ports unique across every route this server owns. + + ``get_open_ports_list`` dedups only within its own call, so we redraw any + port that collides with one already recorded in ``self._allocated_ports`` + (registration/coordinator ports plus every previously-allocated route + port) and register the winners before returning. Callers already hold the + relevant allocation context (``__init__`` is single-threaded; + registration holds ``self._alloc_lock``). + """ + picked: list[int] = [] + # Bounded retry budget so a pathological host (near port exhaustion) + # fails loudly instead of spinning forever. + for _ in range(64): + for port in get_open_ports_list(count=count - len(picked)): + if port in self._allocated_ports: + continue + self._allocated_ports.add(port) + picked.append(port) + if len(picked) == count: + return picked + raise RuntimeError( + f"[OmniMasterServer] Could not allocate {count} unique open ports " + f"(host={self._address}); {len(self._allocated_ports)} already in use." + ) + def _next_free_replica_id(self, stage_id: int) -> int: """Return the next replica id to assign for an auto-assign registration. diff --git a/vllm_omni/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index 7e1153617..a9251a8d8 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -32,7 +32,7 @@ from starlette.routing import Route from starlette.types import ASGIApp, Receive, Scope, Send 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.chat_utils import ChatTemplateConfig, load_chat_template from vllm.entrypoints.launcher import serve_http, terminate_if_errored from vllm.entrypoints.mcp.tool_server import DemoToolServer, MCPToolServer, ToolServer from vllm.entrypoints.openai.api_server import build_app as build_openai_app @@ -60,12 +60,11 @@ 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 from vllm.entrypoints.pooling.scoring.serving import ServingScores -from vllm.entrypoints.serve.disagg.serving import ServingTokens +from vllm.entrypoints.scale_out.token_in_token_out.serving import ServingTokens # vLLM moved `base` from openai.basic.api_router to serve.instrumentator.basic. # Keep a fallback for older/newer upstream layouts during rebase windows. from vllm.entrypoints.serve.instrumentator.basic import base -from vllm.entrypoints.serve.render.serving import OpenAIServingRender from vllm.entrypoints.serve.tokenize.serving import ServingTokenization from vllm.entrypoints.serve.utils.api_utils import ( load_aware_call, @@ -85,6 +84,7 @@ from vllm.entrypoints.speech_to_text.translation.serving import ( OpenAIServingTranslation, ) from vllm.logger import init_logger +from vllm.renderers.online_renderer import OnlineRenderer from vllm.tasks import POOLING_TASKS from vllm.tool_parsers import ToolParserManager from vllm.utils import random_uuid @@ -451,7 +451,7 @@ async def omni_run_server(args, **uvicorn_kwargs) -> None: # Add process-specific prefix to stdout and stderr. decorate_logs("APIServer", skip_if_decorated=True) - listen_address, sock = setup_openai_server(args) + listen_address, sock = setup_openai_server(args, reuse_port=False) # Unified use of omni_run_server_worker, AsyncOmni automatically handles LLM and Diffusion models await omni_run_server_worker(listen_address, sock, args, **uvicorn_kwargs) @@ -830,6 +830,12 @@ async def omni_init_app_state( if tokenizer is None or getattr(tokenizer, "chat_template", None) is None: resolved_chat_template = _load_model_chat_template_json(args.model) + chat_template_config = ChatTemplateConfig( + chat_template=resolved_chat_template, + chat_template_content_format=args.chat_template_content_format, + trust_request_chat_template=args.trust_request_chat_template, + ) + if args.tool_server == "demo": tool_server: ToolServer | None = DemoToolServer() assert isinstance(tool_server, DemoToolServer) @@ -895,11 +901,10 @@ async def omni_init_app_state( await state.openai_serving_models.init_static_loras() # NOTE: kept aligned with upstream `init_app_state`: - # Use OpenAIServingRender (replaces the old OnlineRenderer + OnlineDerenderer + ServingRender split). - state.openai_serving_render = OpenAIServingRender( + # Use OnlineRenderer (replaced OpenAIServingRender which was removed upstream). + state.online_renderer = OnlineRenderer( model_config=engine_client.model_config, renderer=engine_client.renderer, - model_registry=state.openai_serving_models.registry, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -916,7 +921,7 @@ async def omni_init_app_state( OpenAIServingResponses( engine_client, state.openai_serving_models, - openai_serving_render=state.openai_serving_render, + state.online_renderer, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -937,7 +942,7 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, args.response_role, - openai_serving_render=state.openai_serving_render, + online_renderer=state.online_renderer, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -964,7 +969,7 @@ async def omni_init_app_state( OpenAIServingCompletion( engine_client, state.openai_serving_models, - openai_serving_render=state.openai_serving_render, + online_renderer=state.online_renderer, request_logger=request_logger, return_tokens_as_token_ids=args.return_tokens_as_token_ids, enable_prompt_tokens_details=args.enable_prompt_tokens_details, @@ -979,9 +984,7 @@ async def omni_init_app_state( state.openai_serving_models, supported_tasks=tuple(supported_tasks), request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, + chat_template_config=chat_template_config, ) if any(task in POOLING_TASKS for task in supported_tasks) else None @@ -991,9 +994,7 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, + chat_template_config=chat_template_config, ) if "embed" in supported_tasks else None @@ -1003,9 +1004,7 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, request_logger=request_logger, - chat_template=resolved_chat_template, - chat_template_content_format=args.chat_template_content_format, - trust_request_chat_template=args.trust_request_chat_template, + chat_template_config=chat_template_config, ) if "classify" in supported_tasks else None @@ -1014,8 +1013,9 @@ async def omni_init_app_state( ServingScores( engine_client, state.openai_serving_models, + supported_tasks=tuple(supported_tasks), request_logger=request_logger, - score_template=resolved_chat_template, + chat_template_config=chat_template_config, log_error_stack=args.log_error_stack, ) if any(t in supported_tasks for t in ("embed", "score", "token_embed")) @@ -1023,7 +1023,7 @@ async def omni_init_app_state( ) state.serving_tokenization = ServingTokenization( state.openai_serving_models, - state.openai_serving_render, + state.online_renderer, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -1055,7 +1055,7 @@ async def omni_init_app_state( engine_client, state.openai_serving_models, args.response_role, - openai_serving_render=state.openai_serving_render, + online_renderer=state.online_renderer, request_logger=request_logger, chat_template=resolved_chat_template, chat_template_content_format=args.chat_template_content_format, @@ -1074,7 +1074,7 @@ async def omni_init_app_state( ServingTokens( engine_client, state.openai_serving_models, - state.openai_serving_render, + state.online_renderer, 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_audio_generate.py b/vllm_omni/entrypoints/openai/serving_audio_generate.py index 81226439c..f8fedd33f 100644 --- a/vllm_omni/entrypoints/openai/serving_audio_generate.py +++ b/vllm_omni/entrypoints/openai/serving_audio_generate.py @@ -3,7 +3,7 @@ import asyncio import torch from fastapi import Request from fastapi.responses import Response -from vllm.entrypoints.openai.engine.serving import OpenAIServing +from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing from vllm.logger import init_logger from vllm.utils import random_uuid diff --git a/vllm_omni/entrypoints/openai/serving_chat.py b/vllm_omni/entrypoints/openai/serving_chat.py index 2878487fe..6896e3641 100755 --- a/vllm_omni/entrypoints/openai/serving_chat.py +++ b/vllm_omni/entrypoints/openai/serving_chat.py @@ -41,6 +41,7 @@ except ImportError: soundfile = None +from vllm.entrypoints.generate.base.serving import clamp_prompt_logprobs from vllm.entrypoints.launcher import terminate_if_errored from vllm.entrypoints.openai.chat_completion.protocol import ( ChatCompletionNamedToolChoiceParam, @@ -64,7 +65,6 @@ from vllm.entrypoints.openai.engine.protocol import ( ToolCall, UsageInfo, ) -from vllm.entrypoints.openai.engine.serving import clamp_prompt_logprobs from vllm.entrypoints.openai.parser.harmony_utils import ( get_streamable_parser_for_assistant, ) @@ -150,6 +150,21 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): # Harmony flag (always False for vllm-omni models) use_harmony: bool = False + @property + def tool_call_id_type(self) -> str: + """Return the tool call ID type, delegating to model config. + + Upstream vLLM removed the stored ``tool_call_id_type`` attribute + from ``OpenAIServingChat`` after the ParserManager refactor; the + field is now resolved on demand via ``get_tool_call_id_type``. + """ + try: + from vllm.entrypoints.chat_utils import get_tool_call_id_type + + return get_tool_call_id_type(self.model_config) + except Exception: + return "random" + def _should_stream_with_auto_tool_parsing(self, request: ChatCompletionRequest) -> bool: """Check if streamed tokens should go through the tool-call parser. @@ -368,7 +383,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): tool_dicts = [tool.model_dump() for tool in request.tools] if not self.use_harmony: - error_check_ret = self.openai_serving_render.validate_chat_template( + error_check_ret = self.online_renderer.validate_chat_template( request_chat_template=request.chat_template, chat_template_kwargs=request.chat_template_kwargs, trust_request_chat_template=self.trust_request_chat_template, @@ -396,7 +411,7 @@ class OmniOpenAIServingChat(OpenAIServingChat, AudioMixin): ) else: should_include_tools = tool_dicts is not None - conversation, engine_prompts = self.openai_serving_render._make_request_with_harmony( + conversation, engine_prompts = self.online_renderer._make_request_with_harmony( request, should_include_tools ) diff --git a/vllm_omni/entrypoints/openai/serving_speech.py b/vllm_omni/entrypoints/openai/serving_speech.py index 2e9594c02..3db6b63a6 100644 --- a/vllm_omni/entrypoints/openai/serving_speech.py +++ b/vllm_omni/entrypoints/openai/serving_speech.py @@ -21,12 +21,12 @@ import torch from fastapi import HTTPException, Request, UploadFile from fastapi.responses import Response, StreamingResponse from transformers.utils.hub import cached_file +from vllm.entrypoints.generate.base.serving import GenerateBaseServing as OpenAIServing from vllm.entrypoints.launcher import terminate_if_errored from vllm.entrypoints.openai.engine.protocol import ( ErrorResponse, RequestResponseMetadata, ) -from vllm.entrypoints.openai.engine.serving import OpenAIServing from vllm.inputs import tokens_input from vllm.logger import init_logger from vllm.multimodal.media import MediaConnector diff --git a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py index cb84467a3..a39a66b91 100644 --- a/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py +++ b/vllm_omni/experimental/ar_diffusion/kv_cache/paged.py @@ -23,22 +23,7 @@ from dataclasses import dataclass import torch from vllm.v1.core.single_type_kv_cache_manager import SlidingWindowManager from vllm.v1.kv_cache_interface import SlidingWindowSpec - -try: - from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec -except ModuleNotFoundError: - - def register_kv_cache_spec(*, manager_class, uniform_type_base_spec=None): - """Compatibility shim for vLLM versions without kv_cache_spec_registry.""" - - def decorator(spec_class): - from vllm.v1.core.single_type_kv_cache_manager import spec_manager_map - - spec_manager_map[spec_class] = manager_class - return spec_class - - return decorator - +from vllm.v1.kv_cache_spec_registry import register_kv_cache_spec # ── Slot mapping ──────────────────────────────────────────────────────────── diff --git a/vllm_omni/patch.py b/vllm_omni/patch.py index 50eb4d991..4f2441294 100644 --- a/vllm_omni/patch.py +++ b/vllm_omni/patch.py @@ -417,3 +417,70 @@ def _patch_fp8_use_quack_fused_bias(): _patch_fp8_use_quack_fused_bias() + + +# ============================================================================= +# Patch CuMemAllocator._python_free_callback to fix CUDA double-free on shutdown +# ============================================================================= +# WHY: CuMemAllocator._python_free_callback guards the asleep-entry double-free +# skip with ``if data.is_asleep and current_platform.is_rocm():`` — only ROCm +# gets the safe empty-handle return. On CUDA, the callback falls through and +# returns the original handle, causing cuMemRelease on already-freed memory +# (CUDA_ERROR_INVALID_VALUE) during EngineCore subprocess atexit cleanup. +# +# This happens because ``sleep()`` calls ``unmap_and_release()`` on ALL +# platforms, then sets ``is_asleep = True``. When the atexit handler +# (``_shutdown_singleton`` -> ``release_pools()`` -> GC) triggers the free +# callback, the asleep entries must return an empty chunk list so the C +# extension skips ``cuMemRelease`` — exactly the same logic already used by +# the ROCm guard. +# +# The fix removes the ``current_platform.is_rocm()`` condition so the guard +# applies to CUDA (and any future platform that implements cumem). +# +# FRAGILITY: Relies on ``_python_free_callback`` being a regular method +# (not a slot or C extension). If CuMemAllocator is rewritten in C/Cython, +# this monkey-patch will silently become a no-op. +def _patch_cumem_free_callback_cuda() -> None: + try: + from vllm.device_allocator.cumem import CuMemAllocator + except ImportError: + _PATCH_LOGGER.debug("[cumem-cuda] CuMemAllocator not available; skipping patch") + return + + _original_free_callback = CuMemAllocator._python_free_callback + + if getattr(_original_free_callback, "_omni_cumem_cuda_patched", False): + return + + # The upstream bug: `_python_free_callback` only skips the double-free + # for ROCm (line ~206). We wrap the method to extend the guard to all + # platforms. + def _patched_free_callback(self, ptr: int) -> tuple: + data = self.pointer_to_data.pop(ptr) + if data.cpu_backup_tensor is not None: + data.cpu_backup_tensor = None + if data.is_asleep: + # sleep() already called unmap_and_release() on this allocation. + # Return an empty chunk list so the C extension skips + # cuMemRelease, avoiding a double-free. Same logic as the + # existing ROCm guard, but applied to all platforms. + device, size, d_mem, _ = data.handle + result = (device, size, d_mem, []) + _PATCH_LOGGER.debug( + "[cumem-cuda] Free callback: asleep entry %s -> empty handle", + ptr, + ) + return result + # Drain pending kernels before the C extension's cuMemUnmap. + torch.accelerator.synchronize(data.handle[0]) + return data.handle + + _patched_free_callback._omni_cumem_cuda_patched = True + CuMemAllocator._python_free_callback = _patched_free_callback + _PATCH_LOGGER.info( + "[cumem-cuda] CuMemAllocator._python_free_callback patched: asleep guard extended to all platforms." + ) + + +_patch_cumem_free_callback_cuda() diff --git a/vllm_omni/worker/gpu_ar_model_runner.py b/vllm_omni/worker/gpu_ar_model_runner.py index d84404508..6cd2b3566 100644 --- a/vllm_omni/worker/gpu_ar_model_runner.py +++ b/vllm_omni/worker/gpu_ar_model_runner.py @@ -159,6 +159,11 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): async_output_copy_stream = kwargs.pop("async_output_copy_stream") vocab_size = kwargs.pop("vocab_size") routed_experts = kwargs.pop("routed_experts", None) + # Upstream AsyncGPUModelRunnerOutput added check_ep_fault / _has_fault + # for EP all2all fault tolerance (PR #43637). Omni doesn't use this + # feature but must consume the kwarg to prevent TypeError from stray + # kwargs and initialize the attribute so super().get_output() works. + kwargs.pop("check_ep_fault", False) if kwargs: raise TypeError(f"Unexpected OmniAsyncGPUModelRunnerOutput kwargs: {sorted(kwargs)}") @@ -170,6 +175,7 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): self.vocab_size = vocab_size self._logprobs_tensors = logprobs_tensors self._routed_experts = routed_experts + self._has_fault: torch.Tensor | None = None default_stream = torch.cuda.current_stream() with torch.cuda.stream(async_output_copy_stream): @@ -220,6 +226,12 @@ class OmniAsyncGPUModelRunnerOutput(AsyncGPUModelRunnerOutput): if background_exception is not None: raise background_exception self._build_model_runner_output_once() + # Upstream AsyncGPUModelRunnerOutput.get_output() accesses + # self._has_fault for EP all2all fault tolerance (PR #43637). + # Ensure the attribute exists even when __init__ was bypassed + # (e.g. unit tests using object.__new__). + if not hasattr(self, "_has_fault"): + self._has_fault = None with record_function_or_nullcontext("omni_async_output:get_output/finalize_async_sampled_tokens"): return super().get_output()