rebase: align vllm-omni with vLLM 740f379faefd

Signed-off-by: tzhouam <tzhouam@connect.ust.hk>

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():

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

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)

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
 

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,
 )

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

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 ────────────────────────────────────────────────────────────