From b73036f5384f5990fe6e538d15a4803035ac65e9 Mon Sep 17 00:00:00 2001 From: tzhouam Date: Tue, 7 Jul 2026 05:53:51 +0800 Subject: [PATCH] rebase: align vllm-omni with vLLM 740f379faefd Signed-off-by: tzhouam --- tests/ar_diffusion/test_kv_cache.py | 12 ++---------- vllm_omni/engine/async_omni_engine.py | 12 +++++++----- vllm_omni/entrypoints/openai/api_server.py | 2 +- .../openai/serving_audio_generate.py | 2 +- vllm_omni/entrypoints/openai/serving_chat.py | 2 +- vllm_omni/entrypoints/openai/serving_speech.py | 2 +- .../experimental/ar_diffusion/kv_cache/paged.py | 17 +---------------- 7 files changed, 14 insertions(+), 35 deletions(-) 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/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index e71af65d3..8cda746d9 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1567,16 +1567,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/entrypoints/openai/api_server.py b/vllm_omni/entrypoints/openai/api_server.py index f3f0d3584..99f7d925a 100755 --- a/vllm_omni/entrypoints/openai/api_server.py +++ b/vllm_omni/entrypoints/openai/api_server.py @@ -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) 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 43cc14853..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, ) 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 ────────────────────────────────────────────────────────────