chore: import upstream snapshot with attribution

python/ray/llm/_internal/serve/engines/sglang/__init__.py

@@ -0,0 +1,3 @@
+from ray.llm._internal.serve.engines.sglang.sglang_engine import SGLangServer
+
+__all__ = ["SGLangServer"]

python/ray/llm/_internal/serve/engines/sglang/sglang_engine.py

@@ -0,0 +1,773 @@
+"""SGLang engine integration for Ray Serve LLM.
+
+Provides ``SGLangServer``, a custom server class that wraps SGLang's
+in-process engine and exposes chat, completions, embeddings, tokenize,
+and detokenize endpoints through the standard Ray Serve LLM protocol.
+
+Community SGLang support is in early development. Track progress and
+provide feedback at https://github.com/ray-project/ray/issues/61114.
+"""
+
+import copy
+import json
+import signal
+import time
+import uuid
+from typing import (
+    Any,
+    AsyncGenerator,
+    List,
+    Literal,
+    Optional,
+    Union,
+)
+
+from pydantic import BaseModel
+
+from ray.llm._internal.serve.constants import ENABLE_WORKER_PROCESS_SETUP_HOOK
+from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+from ray.llm._internal.serve.core.configs.openai_api_models import (
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    CompletionRequest,
+    CompletionResponse,
+    DetokenizeRequest,
+    DetokenizeResponse,
+    EmbeddingCompletionRequest,
+    EmbeddingRequest,
+    EmbeddingResponse,
+    TokenizeCompletionRequest,
+    TokenizeRequest,
+    TokenizeResponse,
+)
+from ray.llm._internal.serve.core.protocol import RawRequestInfo
+from ray.llm._internal.serve.core.server.llm_server import (
+    _merge_replica_actor_and_child_actor_bundles,
+)
+
+
+class SGLangPauseConfig(BaseModel):
+    """SGLang-specific configuration for pause operation."""
+
+    mode: Literal["abort", "in_place", "retract"] = "abort"
+    """Pause mode:
+    - "abort" (default): Terminate all in-flight requests immediately.
+    - "in_place": Freeze requests in queue, preserve kv cache.
+    - "retract": Freeze requests in queue, free corresponding KV cache.
+    """
+
+
+class SGLangSleepConfig(BaseModel):
+    """SGLang-specific configuration for sleep operation"""
+
+    tags: Optional[List[Literal["kv_cache", "weights", "cuda_graph"]]] = None
+
+    """Sleep tags:
+    - "kv_cache": Discard KV cache
+    - "weights": Offload to CPU RAM
+    - "cuda_graph": Discard CUDA graph
+    - None: Discard/Offload everything
+    """
+
+
+class SGLangWakeupConfig(BaseModel):
+    """SGLang-specific configuration for wakeup operation"""
+
+    tags: Optional[List[Literal["kv_cache", "weights", "cuda_graph"]]] = None
+    """Optional tags to selectively wake up components:
+    - "kv_cache": Restore KV cache only
+    - "weights": Restore weights only
+    - "cuda_graph": Restore CUDA graph only
+    - None: Restore everything
+    """
+
+
+_SLEEP_TAGS: frozenset[str] = frozenset({"kv_cache", "weights", "cuda_graph"})
+
+
+class SGLangServer:
+    def __init__(self, llm_config: LLMConfig):
+
+        self._llm_config = llm_config
+        self.engine_kwargs = llm_config.engine_kwargs
+        self._is_paused = False
+        self._sleeping_tags: set[str] = set()
+
+        try:
+            import sglang
+        except ImportError as e:
+            raise ImportError(
+                "SGLang is not installed or failed to import. Please run "
+                "`pip install sglang[all]` to install required dependencies."
+            ) from e
+
+        # TODO(issue-61108): remove this once sglang#18752 is merged and included
+        # in the minimum supported SGLang version for this example.
+        original_signal_func = signal.signal
+
+        def noop_signal_handler(sig, action):
+            # Returns default handler to satisfy signal.signal() return signature
+            return signal.SIG_DFL
+
+        try:
+            # Override signal.signal with our no-op function
+            signal.signal = noop_signal_handler
+            self.engine = sglang.Engine(**self.engine_kwargs)
+        finally:
+            signal.signal = original_signal_func
+
+    @staticmethod
+    def _build_sampling_params(request: Any) -> dict[str, Any]:
+        sampling_params: dict[str, Any] = {}
+        model_fields_set = getattr(request, "model_fields_set", None)
+        has_model_fields_set = model_fields_set is not None
+        fields_set = set(model_fields_set) if has_model_fields_set else set()
+
+        def was_explicitly_set(field_name: str) -> bool:
+            # Use model_fields_set when available to avoid injecting defaults for
+            # fields omitted by the caller.
+            if has_model_fields_set:
+                return field_name in fields_set
+            return getattr(request, field_name, None) is not None
+
+        temperature = getattr(request, "temperature", None)
+        top_p = getattr(request, "top_p", None)
+        max_tokens = getattr(request, "max_tokens", None)
+        stop = getattr(request, "stop", None)
+
+        if was_explicitly_set("temperature") and temperature is not None:
+            sampling_params["temperature"] = temperature
+        if was_explicitly_set("top_p") and top_p is not None:
+            sampling_params["top_p"] = top_p
+        if was_explicitly_set("max_tokens") and max_tokens is not None:
+            sampling_params["max_new_tokens"] = max_tokens
+        if was_explicitly_set("stop") and stop is not None:
+            sampling_params["stop"] = stop
+
+        return sampling_params
+
+    @staticmethod
+    def _parse_finish_reason(finish_reason_info: Any) -> str:
+        """Parse finish_reason from SGLang metadata."""
+        if isinstance(finish_reason_info, dict):
+            return finish_reason_info.get("type", "length")
+        return str(finish_reason_info)
+
+    @staticmethod
+    def _build_chat_messages(messages: List[Any]) -> List[dict[str, Any]]:
+        converted_messages: List[dict[str, Any]] = []
+        for message in messages:
+            if isinstance(message, dict):
+                message_dict = dict(message)
+            elif hasattr(message, "model_dump") and callable(message.model_dump):
+                message_dict = dict(message.model_dump())
+            else:
+                message_dict = {
+                    "role": getattr(message, "role", "user"),
+                    "content": getattr(message, "content", ""),
+                }
+
+            message_dict["role"] = str(message_dict.get("role", "user"))
+            converted_messages.append(message_dict)
+        return converted_messages
+
+    @staticmethod
+    def _build_chat_template_kwargs(request: Any) -> dict[str, Any]:
+        """
+        Build optional chat-template kwargs using request fields when present.
+        This mirrors SGLang's chat-serving pipeline semantics without directly
+        coupling to its internal server classes.
+
+        Works with both ChatCompletionRequest and TokenizeChatRequest since
+        both expose tools and chat_template_kwargs fields.
+        """
+        kwargs: dict[str, Any] = {}
+
+        tools = getattr(request, "tools", None)
+        if tools is not None:
+            kwargs["tools"] = tools
+
+        reasoning_effort = getattr(request, "reasoning_effort", None)
+        if reasoning_effort is not None:
+            kwargs["reasoning_effort"] = reasoning_effort
+
+        chat_template_kwargs = getattr(request, "chat_template_kwargs", None)
+        if isinstance(chat_template_kwargs, dict):
+            kwargs.update(chat_template_kwargs)
+
+        return kwargs
+
+    def _render_chat_prompt(
+        self,
+        messages: List[dict[str, Any]],
+        add_generation_prompt: bool = True,
+        template_kwargs: Optional[dict[str, Any]] = None,
+    ) -> str:
+        tokenizer = self.engine.tokenizer_manager.tokenizer
+        # SGLang supports --skip-tokenizer-init, where tokenizer is intentionally
+        # None and text prompt rendering is not available.
+        if tokenizer is None:
+            return self._render_fallback_prompt(
+                messages, add_generation_prompt=add_generation_prompt
+            )
+
+        return tokenizer.apply_chat_template(
+            messages,
+            tokenize=False,
+            add_generation_prompt=add_generation_prompt,
+            **(template_kwargs or {}),
+        )
+
+    @staticmethod
+    def _render_fallback_prompt(
+        messages: List[dict[str, Any]],
+        add_generation_prompt: bool = True,
+    ) -> str:
+        # Fallback prompt format for tokenizers without chat-template support.
+        prompt_lines: List[str] = []
+        for message in messages:
+            role = str(message.get("role", "user"))
+            content = message.get("content", "")
+            if content is None:
+                content = ""
+            prompt_lines.append(f"{role}: {content}")
+        if add_generation_prompt:
+            prompt_lines.append("assistant:")
+        return "\n".join(prompt_lines)
+
+    async def start(self) -> None:
+        # Engine is initialized in __init__; keep start idempotent for protocol
+        # compatibility.
+        return
+
+    async def check_health(self) -> None:
+        # SGLang's in-process Engine API does not expose a health-check method.
+        # Its health endpoints exist only in HTTP/gRPC server entrypoints, which
+        # this integration does not run. Keep the protocol hook as a no-op.
+        return
+
+    def _build_generate_kwargs(
+        self, request: Any, prompt: Any, stream: bool
+    ) -> dict[str, Any]:
+        """Build kwargs dict for engine.async_generate."""
+        generate_kwargs: dict[str, Any] = {
+            "prompt": prompt,
+            "stream": stream,
+        }
+        sampling_params = self._build_sampling_params(request)
+        if sampling_params:
+            generate_kwargs["sampling_params"] = sampling_params
+        return generate_kwargs
+
+    async def _generate_raw(
+        self,
+        request: Any,
+        prompt: Any,
+    ) -> dict[str, Any]:
+        """Run generation and return raw engine output payload."""
+        generate_kwargs = self._build_generate_kwargs(request, prompt, stream=False)
+        return await self.engine.async_generate(**generate_kwargs)
+
+    @staticmethod
+    def _extract_generation_metadata(raw: dict[str, Any]) -> dict[str, Any]:
+        """Extract normalized generation metadata from one raw engine payload."""
+        text: str = raw.get("text", "")
+        meta: dict[str, Any] = raw.get("meta_info", {}) or {}
+        finish_reason_info = meta.get("finish_reason", {}) or {}
+        finish_reason = SGLangServer._parse_finish_reason(finish_reason_info)
+
+        prompt_tokens = int(meta.get("prompt_tokens", 0))
+        completion_tokens = int(meta.get("completion_tokens", 0))
+        total_tokens = prompt_tokens + completion_tokens
+
+        return {
+            "text": text.strip(),
+            "id": meta.get("id", f"sglang-gen-{uuid.uuid4().hex}"),
+            "created": int(time.time()),
+            "finish_reason": finish_reason,
+            "prompt_tokens": prompt_tokens,
+            "completion_tokens": completion_tokens,
+            "total_tokens": total_tokens,
+        }
+
+    async def _generate_and_extract_metadata(
+        self,
+        request: Any,
+        prompt: Union[str, List[str]],
+    ) -> Union[dict[str, Any], List[dict[str, Any]]]:
+        """
+        Handles parameter extraction, calls the SGLang engine, and processes the
+        raw response to extract common metadata and generated text.
+
+        Accepts either a single prompt string or a list of prompts. When a list
+        is provided, all prompts are sent to SGLang in one batched call, letting
+        SGLang's scheduler handle concurrency natively via async_generate.
+        """
+        raw = await self._generate_raw(request, prompt)
+
+        # Batch case — SGLang returns a list of results, one per prompt
+        if isinstance(prompt, list):
+            if not raw:
+                raise RuntimeError(
+                    "SGLang engine returned an empty response list during generation."
+                )
+            return [self._extract_generation_metadata(r) for r in raw]
+
+        # Single prompt case
+        if isinstance(raw, list):
+            if not raw:
+                raise RuntimeError(
+                    "SGLang engine returned an empty response list during generation."
+                )
+            raw = raw[0]
+        return self._extract_generation_metadata(raw)
+
+    async def _stream_generate(
+        self,
+        request: Any,
+        prompt: Any,
+    ) -> AsyncGenerator[tuple[str, Optional[str]], None]:
+        """Stream from SGLang engine, yielding (delta_text, finish_reason) tuples.
+
+        SGLang returns cumulative text in each chunk, so this method
+        tracks the previous text and yields only the incremental delta.
+        """
+        generate_kwargs = self._build_generate_kwargs(request, prompt, stream=True)
+        stream = await self.engine.async_generate(**generate_kwargs)
+
+        previous_text = ""
+        async for chunk in stream:
+            text = chunk.get("text", "")
+            meta = chunk.get("meta_info", {}) or {}
+
+            delta_text = text[len(previous_text) :]
+            previous_text = text
+
+            finish_reason_info = meta.get("finish_reason", None)
+            finish_reason = (
+                self._parse_finish_reason(finish_reason_info)
+                if finish_reason_info is not None
+                else None
+            )
+            yield delta_text, finish_reason
+
+    @staticmethod
+    def _build_sse_chunk(
+        gen_id: str,
+        object_type: str,
+        created: int,
+        model: str,
+        choice: dict[str, Any],
+    ) -> str:
+        """Build an SSE-formatted chunk string from a single choice payload."""
+        chunk_data = {
+            "id": gen_id,
+            "object": object_type,
+            "created": created,
+            "model": model,
+            "choices": [choice],
+        }
+        return f"data: {json.dumps(chunk_data)}\n\n"
+
+    async def chat(
+        self,
+        request: ChatCompletionRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[Union[str, ChatCompletionResponse], None]:
+        chat_messages = self._build_chat_messages(request.messages)
+        template_kwargs = self._build_chat_template_kwargs(request)
+        prompt = self._render_chat_prompt(
+            chat_messages, template_kwargs=template_kwargs
+        )
+
+        if request.stream:
+            gen_id = f"sglang-gen-{uuid.uuid4().hex}"
+            created = int(time.time())
+            first_chunk = True
+            async for delta_text, finish_reason in self._stream_generate(
+                request, prompt
+            ):
+                delta: dict[str, Any] = {"content": delta_text}
+                if first_chunk:
+                    delta["role"] = "assistant"
+                    first_chunk = False
+                yield self._build_sse_chunk(
+                    gen_id,
+                    "chat.completion.chunk",
+                    created,
+                    request.model,
+                    {"index": 0, "delta": delta, "finish_reason": finish_reason},
+                )
+            return
+
+        metadata = await self._generate_and_extract_metadata(request, prompt)
+
+        usage_data = {
+            "prompt_tokens": metadata["prompt_tokens"],
+            "completion_tokens": metadata["completion_tokens"],
+            "total_tokens": metadata["total_tokens"],
+        }
+
+        choice_data = {
+            "index": 0,
+            "message": {"role": "assistant", "content": metadata["text"]},
+            "finish_reason": metadata["finish_reason"],
+        }
+
+        resp = ChatCompletionResponse(
+            id=metadata["id"],
+            object="chat.completion",
+            created=metadata["created"],
+            model=request.model,
+            choices=[choice_data],
+            usage=usage_data,
+        )
+
+        yield resp
+
+    async def completions(
+        self,
+        request: CompletionRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[Union[str, CompletionResponse], None]:
+        prompt_input = request.prompt
+
+        # Normalize prompt input.
+        if isinstance(prompt_input, list):
+            if not prompt_input:
+                raise ValueError(
+                    "The 'prompt' list cannot be empty for completion requests."
+                )
+            prompts_to_process = prompt_input
+        else:
+            prompts_to_process = [prompt_input]
+
+        if request.stream:
+            gen_id = f"sglang-gen-{uuid.uuid4().hex}"
+            created = int(time.time())
+            for i, prompt_string in enumerate(prompts_to_process):
+                async for delta_text, finish_reason in self._stream_generate(
+                    request, prompt_string
+                ):
+                    yield self._build_sse_chunk(
+                        gen_id,
+                        "text_completion",
+                        created,
+                        request.model,
+                        {
+                            "index": i,
+                            "text": delta_text,
+                            "logprobs": None,
+                            "finish_reason": finish_reason,
+                        },
+                    )
+            return
+
+        results = await self._generate_and_extract_metadata(request, prompts_to_process)
+
+        all_choices = []
+        total_prompt_tokens = 0
+        total_completion_tokens = 0
+
+        for index, metadata in enumerate(results):
+            total_prompt_tokens += metadata["prompt_tokens"]
+            total_completion_tokens += metadata["completion_tokens"]
+            choice_data = {
+                "index": index,
+                "text": metadata["text"],
+                "logprobs": None,
+                "finish_reason": metadata["finish_reason"],
+            }
+            all_choices.append(choice_data)
+
+        usage_data = {
+            "prompt_tokens": total_prompt_tokens,
+            "completion_tokens": total_completion_tokens,
+            "total_tokens": total_prompt_tokens + total_completion_tokens,
+        }
+
+        last_metadata = results[-1]
+
+        resp = CompletionResponse(
+            id=last_metadata["id"],
+            object="text_completion",
+            created=last_metadata.get("created", int(time.time())),
+            model=getattr(request, "model", "default_model"),
+            choices=all_choices,
+            usage=usage_data,
+        )
+
+        yield resp
+
+    async def embeddings(
+        self,
+        request: EmbeddingRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[EmbeddingResponse, None]:
+        # Input handling follows SGLang's OpenAIServingEmbedding pattern:
+        # https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/openai/serving_embedding.py
+        if isinstance(request, EmbeddingCompletionRequest):
+            prompt = request.input
+        else:
+            # Chat embedding request - join messages without the trailing
+            # "assistant:" generation cue that _render_fallback_prompt adds.
+            chat_messages = self._build_chat_messages(request.messages)
+            prompt = "\n".join(
+                f"{m.get('role', 'user')}: {m.get('content') or ''}"
+                for m in chat_messages
+            )
+
+        # async_encode handles both single strings and lists of strings
+        results = await self.engine.async_encode(prompt)
+        if not isinstance(results, list):
+            results = [results]
+
+        if not results:
+            raise RuntimeError(
+                "SGLang engine returned an empty response for embedding request."
+            )
+
+        # Build response following SGLang's _build_embedding_response pattern
+        data = []
+        total_prompt_tokens = 0
+
+        for idx, ret_item in enumerate(results):
+            data.append(
+                {
+                    "index": idx,
+                    "object": "embedding",
+                    "embedding": ret_item.get("embedding", []),
+                }
+            )
+            meta = ret_item.get("meta_info", {}) or {}
+            total_prompt_tokens += int(meta.get("prompt_tokens", 0))
+
+        resp = EmbeddingResponse(
+            object="list",
+            model=request.model or "",
+            data=data,
+            usage={
+                "prompt_tokens": total_prompt_tokens,
+                "total_tokens": total_prompt_tokens,
+                "completion_tokens": 0,
+            },
+        )
+
+        yield resp
+
+    async def tokenize(
+        self,
+        request: TokenizeRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[TokenizeResponse, None]:
+        tokenizer = self.engine.tokenizer_manager.tokenizer
+        if tokenizer is None:
+            raise RuntimeError(
+                "Tokenizer is not available. The tokenize endpoint is not "
+                "supported when SGLang is initialized with --skip-tokenizer-init."
+            )
+
+        if isinstance(request, TokenizeCompletionRequest):
+            prompt = request.prompt
+        else:
+            # Chat tokenize request - render messages to prompt string
+            chat_messages = self._build_chat_messages(request.messages)
+            add_generation_prompt = getattr(request, "add_generation_prompt", True)
+            template_kwargs = self._build_chat_template_kwargs(request)
+            prompt = self._render_chat_prompt(
+                chat_messages,
+                add_generation_prompt=add_generation_prompt,
+                template_kwargs=template_kwargs,
+            )
+
+        add_special_tokens = getattr(request, "add_special_tokens", True)
+        tokens = tokenizer.encode(prompt, add_special_tokens=add_special_tokens)
+
+        max_model_len = (
+            getattr(self.engine.tokenizer_manager, "context_len", None)
+            or getattr(self.engine.server_args, "context_length", None)
+            or 0
+        )
+
+        yield TokenizeResponse(
+            tokens=tokens,
+            count=len(tokens),
+            max_model_len=max_model_len,
+        )
+
+    async def detokenize(
+        self,
+        request: DetokenizeRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[DetokenizeResponse, None]:
+        tokenizer = self.engine.tokenizer_manager.tokenizer
+        if tokenizer is None:
+            raise RuntimeError(
+                "Tokenizer is not available. The detokenize endpoint is not "
+                "supported when SGLang is initialized with --skip-tokenizer-init."
+            )
+        prompt = tokenizer.decode(request.tokens)
+
+        yield DetokenizeResponse(text=prompt)
+
+    async def llm_config(self) -> Optional[LLMConfig]:
+        return self._llm_config
+
+    @classmethod
+    def get_deployment_options(cls, llm_config: "LLMConfig"):
+        deployment_options = copy.deepcopy(llm_config.deployment_config)
+        pg_config = llm_config.placement_group_config or {}
+
+        ray_actor_options = deployment_options.get("ray_actor_options", {})
+
+        tp_size = llm_config.engine_kwargs.get("tp_size", 1)
+        pp_size = llm_config.engine_kwargs.get("pp_size", 1)
+        num_devices = tp_size * pp_size
+
+        if tp_size < 1 or pp_size < 1:
+            raise ValueError(
+                f"Invalid configuration: tp_size={tp_size} and pp_size={pp_size}. "
+                f"Both must be >= 1."
+            )
+
+        if "placement_group_bundles" not in pg_config:
+            child_bundles = [{"GPU": 1} for _ in range(num_devices)]
+
+            replica_bundle = {
+                "CPU": ray_actor_options.get("num_cpus", 1),
+            }
+
+            if ray_actor_options.get("num_gpus"):
+                replica_bundle["GPU"] = ray_actor_options["num_gpus"]
+
+            replica_bundle.update(ray_actor_options.get("resources", {}))
+
+            if "memory" in ray_actor_options:
+                replica_bundle["memory"] = ray_actor_options["memory"]
+
+            pg_bundles = _merge_replica_actor_and_child_actor_bundles(
+                child_actor_bundles=child_bundles,
+                replica_actor_bundle=replica_bundle,
+            )
+            pg_strategy = "PACK"
+        else:
+            pg_bundles = pg_config.get("placement_group_bundles")
+            pg_strategy = pg_config.get("placement_group_strategy", "PACK")
+
+        deployment_options.update(
+            {
+                "placement_group_bundles": pg_bundles,
+                "placement_group_strategy": pg_strategy,
+            }
+        )
+
+        runtime_env = ray_actor_options.setdefault("runtime_env", {})
+
+        if ENABLE_WORKER_PROCESS_SETUP_HOOK:
+            runtime_env.setdefault(
+                "worker_process_setup_hook",
+                "ray.llm._internal.serve._worker_process_setup_hook",
+            )
+
+        if llm_config.runtime_env:
+            runtime_env.update(llm_config.runtime_env)
+
+        deployment_options["ray_actor_options"] = ray_actor_options
+
+        return deployment_options
+
+    async def pause(self, **kwargs: Any) -> None:
+        """Pause generation on the SGlang server
+
+        This halts generation/encoding requests while keeping model weights in GPU memory. New requests are blocked until resume is called.
+
+        Args:
+            **kwargs: Options parsed into SGLangPauseConfig.
+                - mode (str): "abort" (default), "in_place", or "retract"
+        """
+
+        assert self.engine is not None, "server is not initialized"
+        config = SGLangPauseConfig(**kwargs)
+        from sglang.srt.managers.io_struct import PauseGenerationReqInput
+
+        await self.engine.tokenizer_manager.pause_generation(
+            PauseGenerationReqInput(mode=config.mode)
+        )
+        self._is_paused = True
+
+    async def resume(self, **kwargs: Any) -> None:
+        """Resume generation on the SGLang server after pause.
+
+        Args:
+            **kwargs: Reserved for future options.
+        """
+        assert self.engine is not None, "server is not initialized"
+        from sglang.srt.managers.io_struct import ContinueGenerationReqInput
+
+        await self.engine.tokenizer_manager.continue_generation(
+            ContinueGenerationReqInput()
+        )
+        self._is_paused = False
+
+    async def is_paused(self) -> bool:
+        """Check whether the SGLang server is currently paused.
+
+        Returns:
+            True if the server is paused, False otherwise.
+        """
+        return self._is_paused
+
+    async def sleep(self, **kwargs: Any) -> None:
+        """Put SGLang server to sleep.
+
+        Args:
+            **kwargs: Options parsed into SGLangSleepConfig
+                - tags (List[str], optional): Components to put to sleep.
+        """
+
+        assert self.engine is not None, "server is not initialized"
+        config = SGLangSleepConfig(**kwargs)
+
+        # release_memory_occupation() calls loop.run_until_complete() internally, which fails
+        # inside an async context. Await the underlying coroutine directly.
+        from sglang.srt.entrypoints.engine import ReleaseMemoryOccupationReqInput
+
+        obj = ReleaseMemoryOccupationReqInput(tags=config.tags)
+        await self.engine.tokenizer_manager.release_memory_occupation(obj, None)
+        self._sleeping_tags |= set(config.tags) if config.tags else set(_SLEEP_TAGS)
+
+    async def wakeup(self, **kwargs: Any) -> None:
+        """Wake up the SGLang server from sleep mode.
+
+        Args:
+            **kwargs: Options parsed into SGLangWakeupConfig
+                - tags (List[str], optional): Components to wake up.
+        """
+
+        assert self.engine is not None, "server is not initialized"
+        config = SGLangWakeupConfig(**kwargs)
+        # resume_memory_occupation() release_memory_occupation() calls loop.run_until_complete() internally, which fails
+        # inside an async context. Await the underlying coroutine directly.
+        from sglang.srt.entrypoints.engine import ResumeMemoryOccupationReqInput
+
+        obj = ResumeMemoryOccupationReqInput(tags=config.tags)
+        await self.engine.tokenizer_manager.resume_memory_occupation(obj, None)
+
+        if config.tags is None:
+            self._sleeping_tags.clear()
+        else:
+            self._sleeping_tags -= set(config.tags)
+
+    async def is_sleeping(self) -> bool:
+        """Check whether the SGLang server is currently sleeping.
+
+        Returns:
+            True if any component is currently offloaded/discarded, False otherwise.
+        """
+        return bool(self._sleeping_tags)
+
+    async def reset_prefix_cache(self, timeout: Optional[float] = None) -> None:
+        assert self.engine is not None, "server is not initialized"
+        # flush_cache() calls loop.run_until_complete() internally, which fails
+        # inside an async context. Await the underlying coroutine directly.
+        await self.engine.tokenizer_manager.flush_cache()

python/ray/llm/_internal/serve/engines/vllm/kv_transfer/__init__.py

@@ -0,0 +1,5 @@
+"""KV Transfer connector backends for Ray Serve LLM.
+
+This package provides connector backends for KV cache transfer in vLLM.
+All backends are lazily loaded through the factory to avoid circular imports.
+"""

python/ray/llm/_internal/serve/engines/vllm/kv_transfer/base.py

@@ -0,0 +1,264 @@
+import abc
+import random
+import string
+from typing import TYPE_CHECKING, Any, Dict, Optional, Union
+
+from ray import serve
+
+if TYPE_CHECKING:
+    from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+    from ray.llm._internal.serve.core.configs.openai_api_models import (
+        ChatCompletionRequest,
+        CompletionRequest,
+    )
+
+    # The two OpenAI request models the P/D orchestrator shapes. Defined under
+    # TYPE_CHECKING (and used as a string annotation) to avoid an import cycle
+    # between this module and the config/openai-models modules.
+    RequestType = Union[ChatCompletionRequest, CompletionRequest]
+
+
+def base_prefill_kv_transfer_params() -> Dict[str, Any]:
+    """The ``kv_transfer_params`` common to a prefill (producer) request.
+
+    Tells the prefill engine to produce KV for a remote decode. Connectors layer
+    their own keys (e.g. a transfer id, DP/TP routing) on top of these.
+    """
+    return {
+        "do_remote_decode": True,
+        "do_remote_prefill": False,
+        "remote_engine_id": None,
+        "remote_block_ids": None,
+    }
+
+
+def clamp_request_to_single_token(request: "RequestType") -> None:
+    """Clamp a prefill request to a single, non-streaming token (in place)."""
+    request.max_tokens = 1
+    if hasattr(request, "max_completion_tokens"):
+        request.max_completion_tokens = 1
+    request.stream = False
+    if hasattr(request, "stream_options"):
+        request.stream_options = None
+
+
+class BaseConnectorBackend(abc.ABC):
+    # ---- P/D coordination protocol ----
+    #
+    # These class attributes and methods let the P/D orchestrator
+    # (``PDOrchestratorMixin``) delegate request shaping, peer addressing, and
+    # handoff discipline to the connector. They are connector-agnostic: a
+    # connector picks a quadrant of (``requires_peer_binding``,
+    # ``concurrent_handoff``) and implements ``prepare_prefill_request`` /
+    # ``prepare_decode_request`` accordingly.
+    #
+    # ``requires_peer_binding``:
+    #   * False -> the orchestrator dispatches prefill via the standard handle
+    #     path; the peer (if any) is resolved post-hoc from the prefill response.
+    #   * True  -> the orchestrator selects the prefill replica first
+    #     (``choose_replica``) and passes its ``replica_metadata`` to the backend
+    #     as ``peer`` (pre-dispatch addressing).
+    #
+    # ``concurrent_handoff``:
+    #   * False -> prefill runs to its first chunk before local decode starts
+    #     (sequential handoff).
+    #   * True  -> prefill dispatch and local decode run concurrently.
+    #
+    # The two flags are independent; the known combos:
+    #   * (False, False) — e.g. NixlConnector / LMCacheConnectorV1: decode learns
+    #     everything it needs (remote engine id / block ids) from the prefill
+    #     response, so it must wait for that response (sequential).
+    #   * (True, True)   — e.g. MoRIIO WRITE mode: the peer address is bound
+    #     into the request id before dispatch and prefill *pushes* KV to decode,
+    #     so decode needs nothing from prefill's response and can start
+    #     immediately (concurrent). Concurrent handoff is only possible because
+    #     peer binding happens up front.
+    #   * (True, False)  — e.g. MoRIIO READ mode: the peer is bound up front,
+    #     but decode *pulls* KV using block ids returned in prefill's response,
+    #     so it still waits for prefill to finish (sequential).
+    requires_peer_binding: bool = False
+    concurrent_handoff: bool = False
+
+    def __init__(self, llm_config: "LLMConfig"):
+        """Base class for connector backends.
+
+        Args:
+            llm_config: The llm configuration for this engine
+        """
+        self.llm_config = llm_config
+
+    @property
+    def kv_transfer_config(self) -> Dict[str, Any]:
+        engine_kwargs = self.llm_config.engine_kwargs
+        kv_transfer_config = engine_kwargs.get("kv_transfer_config")
+        assert (
+            kv_transfer_config is not None
+        ), "In Connector backend, kv_transfer_config is not set"
+        return kv_transfer_config
+
+    def _get_unique_suffix(self, len: int = 6) -> str:
+        """Generates unique alphanumeric suffix.
+
+        Args:
+            len: Length of the suffix to generate.
+        Returns:
+            A unique alphanumeric suffix string of specified length.
+        """
+        return "".join(random.choices(string.ascii_letters + string.digits, k=len))
+
+    def _compute_port_offset(self) -> int:
+        """Compute a deterministic port offset for this replica.
+
+        Uses data_parallel_rank if DP case, otherwise falls back to
+        the replica rank assigned by Ray Serve (TP/PP case).
+
+        For TP/PP cases, multiply by num_devices (tp × pp) to reserve
+        sufficient port space, since each worker needs a unique port.
+        Each TP worker adds its tp_rank (0, 1, ..., tp_size-1) to the
+        base port at bind time, and PP stages also need separate ports.
+
+        Returns:
+            Non-negative integer offset to add to a base port.
+        """
+        # Prefer explicit DP rank when available
+        dp_rank = self.llm_config.engine_kwargs.get("data_parallel_rank")
+        if isinstance(dp_rank, int) and dp_rank >= 0:
+            # vLLM already accounts for TP spacing in DP offset calculation
+            # (data_parallel_rank × tp_size), don't multiply here
+            return dp_rank
+
+        # NOTE (jeffreywang): A missing replica context must fail loudly, not
+        # silently return a 0 offset that collides colocated replicas on the
+        # same NIXL side-channel port. get_replica_context() raises RayServeException
+        # outside a replica.
+        rc = serve.get_replica_context()
+        engine_config = self.llm_config.get_engine_config()
+        num_devices = engine_config.num_devices
+        return rc.rank.rank * num_devices
+
+    @abc.abstractmethod
+    def prepare_prefill_request(
+        self, *, request: "RequestType", peer: Optional[Dict[str, Any]]
+    ) -> "RequestType":
+        """Shape the request sent to the remote prefill engine.
+
+        Args:
+            request: The incoming chat/completion request.
+            peer: The selected prefill replica's ``replica_metadata`` dict when
+                the connector opted into pre-dispatch peer binding
+                (``requires_peer_binding=True``), else None.
+
+        Returns:
+            A new request object to dispatch to the prefill engine.
+        """
+        ...
+
+    @abc.abstractmethod
+    def prepare_decode_request(
+        self,
+        *,
+        request: "RequestType",
+        peer: Optional[Dict[str, Any]],
+        prefill_response: Optional[Any],
+    ) -> "RequestType":
+        """Shape the request run on the local decode engine.
+
+        Args:
+            request: The incoming chat/completion request.
+            peer: The selected prefill replica's ``replica_metadata`` dict when
+                the connector opted into pre-dispatch peer binding, else None.
+            prefill_response: The captured prefill response chunk whose
+                ``kv_transfer_params`` may be forwarded, or None when no chunk is
+                captured before decode starts (concurrent-handoff mode).
+
+        Returns:
+            A new request object to run on the local decode engine.
+        """
+        ...
+
+    def setup(self) -> None:
+        """Setup the connector backend.
+
+        This method is called to setup the connector backend.
+        """
+        pass
+
+    def replica_metadata(self) -> Dict[str, Any]:
+        """Static per-replica coordination data published to the orchestrator.
+
+        Surfaced via the replica-metadata hook on ``ReplicaSelection`` so that a
+        connector opting into ``requires_peer_binding`` can address the selected
+        prefill peer. The default backend publishes nothing; connectors that need
+        to advertise an address (e.g. MoRIIO's zmq endpoint) override this.
+
+        Returns:
+            A JSON-serializable dict of per-replica metadata (empty by default).
+        """
+        return {}
+
+
+class DefaultPDProtocolMixin:
+    """The default P/D protocol policy: no peer binding, sequential handoff.
+
+    Implements ``prepare_prefill_request`` / ``prepare_decode_request`` for
+    connectors that follow the standard policy: the prefill engine is told to
+    produce KV for a remote decode (clamped to a single non-streaming token),
+    and the decode engine forwards the ``kv_transfer_params`` that the prefill
+    engine returned on its first response chunk.
+
+    Mix this in *before* ``BaseConnectorBackend`` in a backend's bases so its
+    concrete methods satisfy the abstract methods.
+    """
+
+    def prepare_prefill_request(
+        self, *, request: "RequestType", peer: Optional[Dict[str, Any]]
+    ) -> "RequestType":
+        """Shape the prefill request under the default P/D protocol policy.
+
+        Deep-copies the request, stamps the standard ``kv_transfer_params`` that
+        tell the prefill engine to produce KV for a remote decode, and clamps it
+        to a single, non-streaming token. ``peer`` is ignored.
+        """
+        assert (
+            getattr(request, "kv_transfer_params", None) is None
+        ), "kv_transfer_params should be empty before orchestrator"
+        prefill_request = request.model_copy(deep=True)
+        prefill_request.kv_transfer_params = {
+            **base_prefill_kv_transfer_params(),
+            "remote_host": None,
+            "remote_port": None,
+        }
+        clamp_request_to_single_token(prefill_request)
+        return prefill_request
+
+    def prepare_decode_request(
+        self,
+        *,
+        request: "RequestType",
+        peer: Optional[Dict[str, Any]],
+        prefill_response: Optional[Any],
+    ) -> "RequestType":
+        """Shape the decode request under the default P/D protocol policy.
+
+        Deep-copies the request and, only when a prefill response chunk was
+        captured, forwards its ``kv_transfer_params`` so the decode engine
+        pulls/receives the KV produced by prefill. In concurrent-handoff mode
+        ``prefill_response`` is None and the request is left unmodified. ``peer``
+        is ignored.
+        """
+        decode_request = request.model_copy(deep=True)
+        if prefill_response is not None:
+            decode_request.kv_transfer_params = prefill_response.kv_transfer_params
+        return decode_request
+
+
+class DefaultConnectorBackend(DefaultPDProtocolMixin, BaseConnectorBackend):
+    """Concrete connector backend using the default P/D protocol policy.
+
+    Used as the factory fallback for connectors that are not registered with a
+    dedicated backend class: they get a no-op ``setup()`` and the default
+    request-shaping policy. ``BaseConnectorBackend`` is abstract, so the factory
+    must return a concrete class like this one.
+    """
+
+    pass

python/ray/llm/_internal/serve/engines/vllm/kv_transfer/factory.py

@@ -0,0 +1,140 @@
+"""Factory for lazy-loading KV connector backends.
+
+This module provides a factory pattern for registering and instantiating
+KV connector backends without eagerly importing all implementations.
+This avoids circular import issues and improves startup performance.
+"""
+
+from typing import TYPE_CHECKING, Type, Union
+
+from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+    BaseConnectorBackend,
+    DefaultConnectorBackend,
+)
+from ray.llm._internal.serve.observability.logging import get_logger
+from ray.llm._internal.serve.utils.registry import get_registry
+
+if TYPE_CHECKING:
+    from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+
+
+logger = get_logger(__name__)
+
+# Get the registry instance for KV connector backends
+_kv_backend_registry = get_registry("kv_connector_backend")
+
+
+class KVConnectorBackendFactory:
+    """Factory for creating KV connector backend instances with lazy loading."""
+
+    @classmethod
+    def register_backend(
+        cls,
+        name: str,
+        backend_class_or_path: Union[Type["BaseConnectorBackend"], str],
+    ) -> None:
+        """Register a connector backend.
+
+        This enables the backend to be accessed on every Ray process in the cluster.
+
+        Args:
+            name: The name of the connector (e.g., "LMCacheConnectorV1")
+            backend_class_or_path: Either:
+                - The backend class object directly (preferred), or
+                - A string in the format "module_path:class_name" for lazy loading
+
+        Examples:
+            # Register with class directly (recommended):
+            KVConnectorBackendFactory.register_backend("MyConnector", MyConnectorClass)
+
+            # Register with module path string (for lazy loading):
+            KVConnectorBackendFactory.register_backend("MyConnector", "my.module:MyClass")
+        """
+        _kv_backend_registry.register(name, backend_class_or_path)
+
+    @classmethod
+    def get_backend_class(cls, name: str) -> Type["BaseConnectorBackend"]:
+        """Get the connector backend class by name.
+
+        For registered connectors, returns the registered backend class.
+        For unregistered connectors, returns DefaultConnectorBackend (a concrete
+        backend with a no-op setup() and the default P/D protocol policy),
+        allowing connectors that don't require Ray Serve orchestration to work
+        without registration. (BaseConnectorBackend itself is abstract and
+        cannot be instantiated.)
+
+        Args:
+            name: The name of the connector backend
+
+        Returns:
+            The connector backend class
+
+        Raises:
+            ImportError: If a registered backend fails to load
+        """
+        try:
+            return _kv_backend_registry.get(name)
+        except ValueError:
+            logger.warning(
+                f"Unsupported connector backend: {name}. "
+                f"Using default: {DefaultConnectorBackend.__name__}."
+            )
+            return DefaultConnectorBackend
+        except Exception as e:
+            raise ImportError(
+                f"Failed to load connector backend '{name}': {type(e).__name__}: {e}"
+            ) from e
+
+    @classmethod
+    def create_backend(
+        cls, name: str, llm_config: "LLMConfig"
+    ) -> "BaseConnectorBackend":
+        """Create a connector backend instance.
+
+        Args:
+            name: The name of the connector backend
+            llm_config: The LLM configuration
+
+        Returns:
+            An instance of the connector backend
+        """
+        return cls.get_backend_class(name)(llm_config)
+
+    @classmethod
+    def is_registered(cls, name: str) -> bool:
+        """Check if a connector backend is registered."""
+        return _kv_backend_registry.contains(name)
+
+    @classmethod
+    def unregister_backend(cls, name: str) -> None:
+        """Unregister a connector backend.
+
+        Removes the backend from the registry across all Ray processes.
+
+        Args:
+            name: The name of the connector backend to unregister
+        """
+        _kv_backend_registry.unregister(name)
+
+
+BUILTIN_BACKENDS = {
+    "LMCacheConnectorV1": "ray.llm._internal.serve.engines.vllm.kv_transfer.lmcache:LMCacheConnectorV1Backend",
+    "NixlConnector": "ray.llm._internal.serve.engines.vllm.kv_transfer.nixl:NixlConnectorBackend",
+    "MultiConnector": "ray.llm._internal.serve.engines.vllm.kv_transfer.multi_connector:MultiConnectorBackend",
+    "MoRIIOConnector": "ray.llm._internal.serve.engines.vllm.kv_transfer.moriio:MoRIIOConnectorBackend",
+}
+
+
+def _initialize_registry() -> None:
+    """Initialize the registry with built-in backends.
+
+    This function is called when the module is imported to ensure
+    built-in backends are registered.
+    """
+    for name, backend_path in BUILTIN_BACKENDS.items():
+        if not KVConnectorBackendFactory.is_registered(name):
+            KVConnectorBackendFactory.register_backend(name, backend_path)
+
+
+# Initialize registry when module is imported
+_initialize_registry()

python/ray/llm/_internal/serve/engines/vllm/kv_transfer/lmcache.py

@@ -0,0 +1,62 @@
+from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+    BaseConnectorBackend,
+    DefaultPDProtocolMixin,
+)
+from ray.llm._internal.serve.observability.logging import get_logger
+
+logger = get_logger(__name__)
+
+
+def _check_lmcache_installed():
+    try:
+        import lmcache  # noqa: F401
+    except ImportError:
+        raise ImportError(
+            "LMCache is not installed. Please install it with `pip install lmcache`."
+        )
+
+
+class LMCacheConnectorV1Backend(DefaultPDProtocolMixin, BaseConnectorBackend):
+
+    KV_CONNECTOR_EXTRA_CONFIG_FIELD_NAME = "kv_connector_extra_config"
+    LMCACHE_RPC_PORT_FIELD_NAME = "lmcache_rpc_port"
+    DEFAULT_LMCACHE_RPC_PORT_NAME = "lmcache_rpc_port"
+
+    def setup(self) -> None:
+        """Initialize the LMCache connector backend.
+
+        Creates a unique LMCache RPC port name across replicas by appending
+        a random suffix to the base port name.
+
+        Raises:
+            ImportError: If LMCache is not installed.
+        """
+        _check_lmcache_installed()
+
+        if (
+            LMCacheConnectorV1Backend.KV_CONNECTOR_EXTRA_CONFIG_FIELD_NAME
+            not in self.kv_transfer_config
+        ):
+            return
+
+        kv_connector_extra_config = self.kv_transfer_config[
+            LMCacheConnectorV1Backend.KV_CONNECTOR_EXTRA_CONFIG_FIELD_NAME
+        ]
+        base_value = kv_connector_extra_config.get(
+            LMCacheConnectorV1Backend.LMCACHE_RPC_PORT_FIELD_NAME,
+            LMCacheConnectorV1Backend.DEFAULT_LMCACHE_RPC_PORT_NAME,
+        )
+
+        # Append random suffix for uniqueness
+        lmcache_rpc_port_value = str(base_value) + self._get_unique_suffix()
+        if (
+            LMCacheConnectorV1Backend.LMCACHE_RPC_PORT_FIELD_NAME
+            in kv_connector_extra_config
+        ):
+            logger.info(
+                f"Setting unique lmcache_rpc_port={lmcache_rpc_port_value} for current replica."
+            )
+
+        kv_connector_extra_config[
+            LMCacheConnectorV1Backend.LMCACHE_RPC_PORT_FIELD_NAME
+        ] = lmcache_rpc_port_value

python/ray/llm/_internal/serve/engines/vllm/kv_transfer/moriio.py

@@ -0,0 +1,344 @@
+"""MoRIIO connector backend for Ray Serve LLM (analogue of nixl.py).
+
+Configures a vLLM engine's ``kv_transfer_config.kv_connector_extra_config`` for
+the MoRIIO connector and computes per-replica handshake/notify ports so colocated
+replicas don't collide. Also builds the engine's advertised zmq address so the
+P/D orchestrator can discover it via the replica-metadata hook
+(``ReplicaSelection.replica_metadata``), and implements the PD connector protocol
+(``requires_peer_binding`` / ``concurrent_handoff`` / ``prepare_prefill_request`` /
+``prepare_decode_request``) so the decode orchestrator can address the selected
+prefill peer by request id.
+
+Unlike NIXL/LMCache, MoRIIO does NOT use ``DefaultPDProtocolMixin``: it has custom
+request shaping (a dual-address request_id + transfer_id) and therefore IMPLEMENTS
+the abstract ``prepare_*`` methods directly on ``BaseConnectorBackend``.
+
+Two transfer disciplines, selected by ``read_mode``:
+  * WRITE (default): prefill PUSHES KV to decode -> concurrent handoff.
+  * READ: decode PULLS KV from prefill -> sequential handoff; the decode request
+    forwards the ``remote_block_ids`` / ``remote_engine_id`` the prefill engine
+    returned.
+
+The dual-address request_id and the transfer_id are derived DETERMINISTICALLY
+from the incoming request id (a hash), so ``prepare_prefill_request`` and
+``prepare_decode_request`` produce identical ids across their two separate calls
+without per-request backend state (the backend instance is shared across
+requests).
+
+Registered with Ray's public connector registry via the factory.
+"""
+
+import hashlib
+import logging
+import re
+from typing import TYPE_CHECKING, Any, Dict, Optional, Tuple
+
+import ray
+from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+    BaseConnectorBackend,
+    base_prefill_kv_transfer_params,
+    clamp_request_to_single_token,
+)
+
+if TYPE_CHECKING:
+    from ray.llm._internal.serve.engines.vllm.kv_transfer.base import RequestType
+
+logger = logging.getLogger(__name__)
+
+# Defaults mirror vLLM's MoRIIOConstants (DEFAULT_HANDSHAKE_PORT / NOTIFY_PORT).
+# Prefill uses these bases; decode is shifted (see builder.py) so a colocated
+# P+D pair on one node doesn't collide.
+DEFAULT_HANDSHAKE_PORT_BASE = 6301
+DEFAULT_NOTIFY_PORT_BASE = 61005
+
+# experimental_configs keys understood by this backend.
+HANDSHAKE_PORT_BASE_KEY = "MORI_HANDSHAKE_PORT_BASE"
+NOTIFY_PORT_BASE_KEY = "MORI_NOTIFY_PORT_BASE"
+
+# ---------------------------------------------------------------------------
+# Dual-address request_id / zmq address encoding.
+#
+# These MUST stay byte-compatible with the regexes vLLM's MoRIIO connector uses
+# to recover peer addresses from the request_id:
+#
+#   vllm/distributed/kv_transfer/kv_connector/v1/moriio/moriio_common.py
+#       _PREFILL_ZMQ_RE = re.compile(r"___prefill_addr_(.+?)___decode_addr_")
+#       _DECODE_ZMQ_RE  = re.compile(r"___decode_addr_(.+)_[0-9a-f]{32}(?:-.*)?$")
+#       # zmq address: "host:IP,handshake:PORT,notify:PORT"
+# ---------------------------------------------------------------------------
+
+_PREFILL_PREFIX = "___prefill_addr_"
+_DECODE_PREFIX = "___decode_addr_"
+_TRANSFER_PREFIX = "tx"
+
+# Copies of vLLM's regexes for local validation / round-trip tests.
+_PREFILL_ZMQ_RE = re.compile(r"___prefill_addr_(.+?)___decode_addr_")
+_DECODE_ZMQ_RE = re.compile(r"___decode_addr_(.+)_[0-9a-f]{32}(?:-.*)?$")
+
+
+def build_zmq_address(host: str, handshake_port: int, notify_port: int) -> str:
+    """Build the MORI zmq address string ``host:IP,handshake:PORT,notify:PORT``."""
+    return f"host:{host},handshake:{handshake_port},notify:{notify_port}"
+
+
+def parse_zmq_address(zmq_address: str) -> Tuple[str, int, int]:
+    """Inverse of :func:`build_zmq_address` -> ``(host, handshake_port, notify_port)``."""
+    parts = {}
+    for segment in zmq_address.split(","):
+        key, _, val = segment.partition(":")
+        parts[key.strip()] = val.strip()
+    return parts["host"], int(parts["handshake"]), int(parts["notify"])
+
+
+def parse_peer_zmq(request_id: str, is_producer: bool) -> str:
+    """Recover the peer's zmq address from a request id (for tests/debugging).
+
+    Producer (prefill) wants the *decode* address; consumer wants the *prefill*.
+    """
+    rex = _DECODE_ZMQ_RE if is_producer else _PREFILL_ZMQ_RE
+    m = rex.search(request_id)
+    if not m:
+        raise ValueError(f"No peer zmq address in request_id: {request_id!r}")
+    return m.group(1)
+
+
+def _read_mode_enabled(extra_config: Dict[str, Any]) -> bool:
+    """Mirror vLLM's ``get_moriio_mode`` parse of ``read_mode``.
+
+    true / 1 -> READ; anything else -> WRITE (default).
+    """
+    return str(extra_config.get("read_mode", "false")).lower().strip() in (
+        "true",
+        "1",
+    )
+
+
+class MoRIIOConnectorBackend(BaseConnectorBackend):
+    """Set up MoRIIO ports/extra_config and implement the PD connector protocol."""
+
+    # The advertised zmq address ("host:IP,handshake:PORT,notify:PORT"),
+    # computed by setup(); consumers reach it via this backend instance.
+    _zmq_address: Optional[str] = None
+
+    # MORI addresses peers by the dual-address request id, so the orchestrator
+    # must bind to the selected prefill replica BEFORE dispatch.
+    requires_peer_binding: bool = True
+
+    def _extra_config(self) -> dict:
+        cfg = self.kv_transfer_config.setdefault("kv_connector_extra_config", {})
+        return cfg
+
+    @property
+    def _read_mode(self) -> bool:
+        """True iff this engine's MoRIIO connector is configured for READ mode."""
+        extra = self._extra_config()
+        return _read_mode_enabled(extra)
+
+    @property
+    def concurrent_handoff(self) -> bool:
+        """WRITE -> concurrent (prefill pushes); READ -> sequential (decode pulls)."""
+        return not self._read_mode
+
+    def setup(self) -> None:
+        offset = self._compute_port_offset()
+
+        handshake_base = int(
+            self.llm_config.experimental_configs.get(
+                HANDSHAKE_PORT_BASE_KEY, DEFAULT_HANDSHAKE_PORT_BASE
+            )
+        )
+        notify_base = int(
+            self.llm_config.experimental_configs.get(
+                NOTIFY_PORT_BASE_KEY, DEFAULT_NOTIFY_PORT_BASE
+            )
+        )
+
+        # NOTE: vLLM internally adds get_port_offset(dp_rank, tp_rank) on top of
+        # these bases. For TP/DP>1, reserve a stride >= tp_size*pp_size when
+        # shifting decode's base in the builder so the two offset schemes never
+        # overlap.
+        handshake_port = handshake_base + offset
+        notify_port = notify_base + offset
+
+        extra = self._extra_config()
+        # Required keys for vLLM's config parser (KeyError otherwise) -- proxyless.
+        extra.setdefault("proxy_ip", "")  # empty => ping/registration thread disabled
+        extra.setdefault("proxy_ping_port", "0")
+        # TODO: real Serve replica HTTP port. Harmless placeholder while
+        # proxy_ip="" (only used to build request_address for the disabled ping).
+        extra.setdefault("http_port", str(8000 + offset))
+        # WRITE mode (prefill pushes). READ would be "true".
+        extra.setdefault("read_mode", "false")
+        extra["handshake_port"] = str(handshake_port)
+        extra["notify_port"] = str(notify_port)
+
+        # Advertise the Ray internal cluster IP as the zmq host.
+        host = ray.util.get_node_ip_address()
+        zmq_address = build_zmq_address(host, handshake_port, notify_port)
+        # Stash so replica_metadata() can publish it; the decode
+        # orchestrator reads the selected prefill replica's copy off the peer.
+        self._zmq_address = zmq_address
+        # Cross-node correctness: vLLM's MoRIIO worker otherwise binds/advertises
+        # get_ip(), which on a Ray cluster is the node's unroutable public IP, and
+        # VLLM_HOST_IP cannot be propagated to workers (it is excluded from vLLM's
+        # driver->worker env copy). Pass the routable node IP through the connector
+        # config, which does reach the workers; vLLM's MoRIIO connector honors
+        # "host_ip" over get_ip(). Requires vllm-project/vllm#45488.
+        extra["host_ip"] = host
+
+    # ---- parallelism (data/tensor) ----
+
+    def _dp_rank(self) -> int:
+        rank = self.llm_config.engine_kwargs.get("data_parallel_rank")
+        return rank if isinstance(rank, int) and rank >= 0 else 0
+
+    def _dp_size(self) -> int:
+        return int(self.llm_config.engine_kwargs.get("data_parallel_size") or 1)
+
+    def _tp_size(self) -> int:
+        return int(self.llm_config.engine_kwargs.get("tensor_parallel_size") or 1)
+
+    # ---- replica metadata (published via the replica-metadata hook) ----
+
+    def replica_metadata(self) -> dict:
+        """Static per-replica coordination data published to the orchestrator.
+
+        The prefill replica publishes its MORI zmq address and its parallelism
+        (DP rank/size, TP size); the decode orchestrator reads them off the
+        selected prefill replica's ``ReplicaSelection.replica_metadata`` and uses
+        them to address the right remote (dp_rank, tp) workers.
+        """
+        return {
+            "mori_zmq_address": self._zmq_address,
+            "dp_rank": self._dp_rank(),
+            "dp_size": self._dp_size(),
+            "tp_size": self._tp_size(),
+        }
+
+    def _remote_routing(self, remote: Dict[str, Any]) -> Dict[str, Any]:
+        """``kv_transfer_params`` keys telling vLLM which remote workers to reach.
+
+        ``remote`` is the metadata of the *other* side of the transfer: the
+        decode (this orchestrator) for a prefill request, the selected prefill
+        peer for a decode request. vLLM addresses a remote worker at
+        ``advertised_base + get_port_offset(remote_dp_rank, tp_index)`` and
+        handshakes all ``remote_dp_size`` ranks, so both must match the target
+        replica. ``tp_size`` is the remote's TP (symmetric across P/D).
+        """
+        return {
+            "remote_dp_rank": int(remote.get("dp_rank", 0)),
+            "remote_dp_size": int(remote.get("dp_size", 1)),
+            "tp_size": int(remote.get("tp_size", self._tp_size())),
+        }
+
+    def _own_routing(self) -> Dict[str, Any]:
+        """``_remote_routing`` input describing this (decode orchestrator) replica
+        -- the remote side for a prefill request."""
+        return {
+            "dp_rank": self._dp_rank(),
+            "dp_size": self._dp_size(),
+            "tp_size": self._tp_size(),
+        }
+
+    # ---- request shaping (PD connector protocol) ----
+
+    def _dual_ids(
+        self, request: Any, peer: Optional[Dict[str, Any]]
+    ) -> Tuple[str, str]:
+        """Compute the (dual-address request_id, transfer_id) for this request.
+
+        ``prepare_prefill_request`` and ``prepare_decode_request`` are two
+        independent, stateless calls for the same request, so both ids are
+        derived deterministically (hash of a stable per-request seed) — no
+        per-request backend state.
+        """
+        prefill_zmq = (peer or {}).get("mori_zmq_address")
+        decode_zmq = self._zmq_address
+        if not prefill_zmq:
+            raise ValueError(
+                "MoRIIO peer is missing 'mori_zmq_address': the selected prefill "
+                "replica did not publish its address (is MoRIIOConnector "
+                "configured on the prefill deployment?)."
+            )
+        if not decode_zmq:
+            raise ValueError(
+                "MoRIIO decode zmq address is not set: setup() must run on this "
+                "engine before requests are shaped."
+            )
+        # The incoming request_id (always populated -- OpenAI models default it
+        # to a uuid) is the seed. Both prepare_* calls run on the same request
+        # object, so they agree; uniqueness per request is inherited from it.
+        seed = str(request.request_id)
+        # 32 hex chars (the trailing uid _PREFILL_ZMQ_RE / _DECODE_ZMQ_RE
+        # anchor on); a hash of the seed, so both prepare_* calls agree.
+        uid = hashlib.sha256(seed.encode()).hexdigest()[:32]
+        # Wire format consumed by vLLM's MoRIIO connector.
+        request_id = f"{_PREFILL_PREFIX}{prefill_zmq}{_DECODE_PREFIX}{decode_zmq}_{uid}"
+        transfer_id = f"{_TRANSFER_PREFIX}-{uid}"
+        return request_id, transfer_id
+
+    def prepare_prefill_request(
+        self, *, request: "RequestType", peer: Optional[Dict[str, Any]]
+    ) -> "RequestType":
+        request_id, transfer_id = self._dual_ids(request, peer)
+        prefill_request = request.model_copy(deep=True)
+        # The dual-address id (peer zmq encoded in it) must reach the engine:
+        # setting request_id explicitly makes the LLMServer pipeline preserve it
+        # (not clobber it with the Serve id) and the engine copies it into the
+        # X-Request-Id header that vLLM's MoRIIO connector parses.
+        prefill_request.request_id = request_id
+        prefill_request.kv_transfer_params = {
+            **base_prefill_kv_transfer_params(),
+            "transfer_id": transfer_id,
+            # The prefill engine's remote is the decode (this orchestrator).
+            **self._remote_routing(self._own_routing()),
+        }
+        clamp_request_to_single_token(prefill_request)
+        return prefill_request
+
+    def prepare_decode_request(
+        self,
+        *,
+        request: "RequestType",
+        peer: Optional[Dict[str, Any]],
+        prefill_response: Optional[Any],
+    ) -> "RequestType":
+        request_id, transfer_id = self._dual_ids(request, peer)
+        decode_request = request.model_copy(deep=True)
+        decode_request.request_id = request_id
+        # The decode engine's remote is the selected prefill peer.
+        remote_routing = self._remote_routing(peer or {})
+
+        if not self._read_mode:
+            # WRITE: prefill pushes KV; decode just needs do_remote_prefill + the
+            # shared transfer_id (no block ids -- they are pushed, not pulled).
+            decode_request.kv_transfer_params = {
+                "do_remote_prefill": True,
+                "do_remote_decode": False,
+                "remote_engine_id": None,
+                "remote_block_ids": None,
+                "transfer_id": transfer_id,
+                **remote_routing,
+            }
+            return decode_request
+
+        # READ: decode PULLS KV; forward the remote_block_ids / remote_engine_id
+        # the prefill engine returned on its response. If absent (e.g. prompt <
+        # block_size / full prefix hit), fall back to a local recompute.
+        prefill_kv_params = getattr(prefill_response, "kv_transfer_params", None)
+        params = dict(prefill_kv_params) if prefill_kv_params else {}
+        if params.get("remote_block_ids") and params.get("remote_engine_id"):
+            params.setdefault("transfer_id", transfer_id)
+            params["do_remote_prefill"] = True
+            params["do_remote_decode"] = False
+            # Address the prefill peer's (dp_rank, dp_size, tp) workers.
+            params.update(remote_routing)
+            decode_request.kv_transfer_params = params
+        else:
+            logger.warning(
+                "[MORI][READ] prefill returned no remote_block_ids/remote_engine_id "
+                "(kv_transfer_params=%s); decode will recompute locally.",
+                prefill_kv_params,
+            )
+            decode_request.kv_transfer_params = None
+        return decode_request

python/ray/llm/_internal/serve/engines/vllm/kv_transfer/multi_connector.py

@@ -0,0 +1,94 @@
+import copy
+from typing import TYPE_CHECKING, List
+
+from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+    BaseConnectorBackend,
+)
+from ray.llm._internal.serve.engines.vllm.kv_transfer.factory import (
+    KVConnectorBackendFactory,
+)
+
+if TYPE_CHECKING:
+    from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+
+
+class MultiConnectorBackend(BaseConnectorBackend):
+    """Wraps multiple sub-connectors.
+
+    The P/D protocol (``prepare_prefill_request`` / ``prepare_decode_request`` and
+    the ``requires_peer_binding`` / ``concurrent_handoff`` policy) is delegated to
+    the *first* (top-most) sub-connector listed in ``connectors`` — that
+    connector's policy governs request shaping and handoff for the group. Each
+    sub-connector's ``setup()`` still runs.
+    """
+
+    def __init__(self, llm_config: "LLMConfig"):
+        super().__init__(llm_config)
+        self._connector_backends: List[BaseConnectorBackend] = []
+
+    def setup(self) -> None:
+        """Setup all connectors listed in the kv_transfer_config."""
+        kv_transfer_config = self.kv_transfer_config
+        connectors = kv_transfer_config.get("kv_connector_extra_config", {}).get(
+            "connectors", []
+        )
+        if not connectors:
+            # Fail fast at setup rather than with a cryptic error when the
+            # orchestrator later delegates to a (missing) top-most sub-connector.
+            raise ValueError(
+                "MultiConnector requires at least one sub-connector in "
+                "kv_connector_extra_config.connectors."
+            )
+
+        for connector in connectors:
+            connector_backend_str = connector.get("kv_connector")
+            if connector_backend_str is None:
+                raise ValueError("kv_connector is not set in the connector")
+
+            if connector_backend_str == "MultiConnector":
+                raise ValueError(
+                    "Nesting MultiConnector within MultiConnector is not supported."
+                )
+
+            # Merge parent config with connector-specific config
+            sub_llm_config = copy.deepcopy(self.llm_config)
+            sub_llm_config.engine_kwargs["kv_transfer_config"] = {
+                **{
+                    k: v
+                    for k, v in kv_transfer_config.items()
+                    if k != "kv_connector_extra_config"
+                },
+                **connector,
+            }
+
+            # Use factory to get backend class lazily
+            connector_backend = KVConnectorBackendFactory.create_backend(
+                connector_backend_str, sub_llm_config
+            )
+            connector_backend.setup()
+            self._connector_backends.append(connector_backend)
+
+    @property
+    def _primary(self) -> BaseConnectorBackend:
+        """The top-most sub-connector, whose protocol governs the group."""
+        if not self._connector_backends:
+            raise ValueError(
+                "MultiConnectorBackend has no sub-connectors; was setup() called?"
+            )
+        return self._connector_backends[0]
+
+    @property
+    def requires_peer_binding(self) -> bool:
+        return bool(self._connector_backends) and self._primary.requires_peer_binding
+
+    @property
+    def concurrent_handoff(self) -> bool:
+        return bool(self._connector_backends) and self._primary.concurrent_handoff
+
+    def prepare_prefill_request(self, *, request, peer):
+        return self._primary.prepare_prefill_request(request=request, peer=peer)
+
+    def prepare_decode_request(self, *, request, peer, prefill_response):
+        return self._primary.prepare_decode_request(
+            request=request, peer=peer, prefill_response=prefill_response
+        )

python/ray/llm/_internal/serve/engines/vllm/kv_transfer/nixl.py

@@ -0,0 +1,67 @@
+import os
+
+import ray
+from ray.llm._internal.serve.engines.vllm.kv_transfer.base import (
+    BaseConnectorBackend,
+    DefaultPDProtocolMixin,
+)
+
+
+class NixlConnectorBackend(DefaultPDProtocolMixin, BaseConnectorBackend):
+    def _set_side_channel_port(self):
+        from vllm import envs as vllm_envs
+
+        if vllm_envs.is_set("VLLM_NIXL_SIDE_CHANNEL_PORT"):
+            return
+
+        base_port = int(
+            self.llm_config.experimental_configs.get(
+                "NIXL_SIDE_CHANNEL_PORT_BASE", 20000
+            )
+        )
+        port = base_port + self._compute_port_offset()
+        os.environ["VLLM_NIXL_SIDE_CHANNEL_PORT"] = str(port)
+
+    def _set_side_channel_host(self):
+        from vllm import envs as vllm_envs
+
+        if not vllm_envs.is_set("VLLM_NIXL_SIDE_CHANNEL_HOST"):
+            # Use Ray's node IP (internal/cluster IP) instead of vLLM's
+            # get_ip() which can return external/public IPs on hostNetwork
+            # pods, causing cross-node NIXL handshakes to fail.
+            os.environ["VLLM_NIXL_SIDE_CHANNEL_HOST"] = ray.util.get_node_ip_address()
+
+    def setup(self) -> None:
+        """Initialize the NIXL connector backend.
+
+        This method sets up the NIXL (NVIDIA Inference Xfer Library) connector by:
+        1. Verifying that the required vLLM environment variables are supported
+        2. Configuring the side channel port and host if not already set
+        3. Creating a unique engine ID across replicas
+
+        The side channel is used for KV cache transfer between vLLM instances.
+
+        Raises:
+            ValueError: If the current vLLM version doesn't support the required
+                       NIXL environment variables.
+        """
+        from vllm import envs as vllm_envs
+
+        if (
+            "VLLM_NIXL_SIDE_CHANNEL_PORT" not in vllm_envs.environment_variables
+            or "VLLM_NIXL_SIDE_CHANNEL_HOST" not in vllm_envs.environment_variables
+        ):
+            raise ValueError(
+                "This vLLM version does not support VLLM_NIXL_SIDE_CHANNEL_PORT"
+                "or VLLM_NIXL_SIDE_CHANNEL_HOST environment variable. It's likely"
+                "that you are using an older version of vLLM."
+            )
+
+        self._set_side_channel_port()
+        self._set_side_channel_host()
+
+        # We need to overwrite the engine_id to make it unique across replicas.
+        engine_id = self.kv_transfer_config.get("engine_id", self._get_unique_suffix())
+        host = vllm_envs.VLLM_NIXL_SIDE_CHANNEL_HOST
+        port = vllm_envs.VLLM_NIXL_SIDE_CHANNEL_PORT
+        self.kv_transfer_config["engine_id"] = "-".join([engine_id, host, str(port)])

python/ray/llm/_internal/serve/engines/vllm/vllm_engine.py

@@ -0,0 +1,960 @@
+import argparse
+import dataclasses
+import inspect
+import json
+import typing
+from typing import (
+    TYPE_CHECKING,
+    Any,
+    AsyncGenerator,
+    Dict,
+    List,
+    Literal,
+    Optional,
+    Tuple,
+    Union,
+)
+
+from pydantic import BaseModel, field_validator
+from starlette.datastructures import State
+from starlette.requests import Request
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.entrypoints.openai.cli_args import FrontendArgs
+from vllm.entrypoints.openai.engine.protocol import ErrorResponse as VLLMErrorResponse
+
+import ray
+from ray.llm._internal.common.callbacks.base import CallbackCtx
+from ray.llm._internal.common.utils.import_utils import try_import
+from ray.llm._internal.serve.core.configs.llm_config import (
+    DiskMultiplexConfig,
+    LLMConfig,
+)
+from ray.llm._internal.serve.core.configs.openai_api_models import (
+    ChatCompletionRequest,
+    ChatCompletionResponse,
+    CompletionRequest,
+    CompletionResponse,
+    DetokenizeRequest,
+    DetokenizeResponse,
+    EmbeddingRequest,
+    EmbeddingResponse,
+    ErrorInfo,
+    ErrorResponse,
+    ScoreRequest,
+    ScoreResponse,
+    TokenizeRequest,
+    TokenizeResponse,
+    TranscriptionRequest,
+    TranscriptionResponse,
+)
+from ray.llm._internal.serve.core.engine.protocol import LLMEngine
+from ray.llm._internal.serve.core.protocol import RawRequestInfo
+from ray.llm._internal.serve.engines.vllm.vllm_models import (
+    VLLMEngineConfig,
+)
+from ray.llm._internal.serve.observability.logging import get_logger
+from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_router import (
+    is_kv_aware,
+)
+from ray.llm._internal.serve.routing_policies.kv_aware.vllm.kv_events import (
+    assign_replica_kv_events_endpoint,
+    get_kv_event_routing_stats,
+)
+from ray.llm._internal.serve.routing_policies.kv_aware.vllm.token_tracking import (
+    enable_token_tracking,
+)
+from ray.llm._internal.serve.utils.node_initialization_utils import (
+    initialize_node,
+)
+from ray.util.placement_group import PlacementGroup
+from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+    from vllm.engine.protocol import EngineClient
+    from vllm.entrypoints.openai.chat_completion.serving import OpenAIServingChat
+    from vllm.entrypoints.openai.completion.serving import OpenAIServingCompletion
+    from vllm.entrypoints.openai.models.serving import OpenAIServingModels
+    from vllm.entrypoints.pooling.embed.serving import ServingEmbedding
+    from vllm.entrypoints.pooling.scoring.serving import ServingScores
+    from vllm.entrypoints.serve.tokenize.serving import OpenAIServingTokenization
+    from vllm.entrypoints.speech_to_text.transcription.serving import (
+        OpenAIServingTranscription,
+    )
+
+vllm = try_import("vllm")
+logger = get_logger(__name__)
+
+
+def _canonicalize_request_id_header(
+    request: Any, raw_request_info: Optional[RawRequestInfo]
+) -> Optional[RawRequestInfo]:
+    """Ensure raw_request_info carries X-Request-Id == request.request_id so vLLM's
+    OpenAI layer (which prefers the header) sees the authoritative request id.
+
+    Returns a RawRequestInfo (new or mutated) with a single, correctly-cased header.
+    If the request has no request_id, this is a no-op and returns raw_request_info
+    unchanged (so embeddings/score/transcription requests are unaffected).
+    """
+    rid = getattr(request, "request_id", None)
+    if not rid:
+        return raw_request_info
+    headers = dict(raw_request_info.headers) if raw_request_info is not None else {}
+    # Drop any existing variant of the header (case- and separator-insensitive,
+    # e.g. "X-Request-Id" or "x_request_id") before setting the canonical one.
+    headers = {
+        k: v
+        for k, v in headers.items()
+        if k.replace("_", "-").lower() != "x-request-id"
+    }
+    headers["x-request-id"] = str(rid)
+    if raw_request_info is None:
+        return RawRequestInfo(headers=headers)
+    # Preserve any non-header fields RawRequestInfo carries (now or in the future).
+    return dataclasses.replace(raw_request_info, headers=headers)
+
+
+def _convert_config_dicts(merged: dict) -> dict:
+    """Convert dict values to their proper vLLM config classes based on type hints.
+
+    vLLM's AsyncEngineArgs has fields like structured_outputs_config,
+    compilation_config, etc. that expect dataclass instances. When users pass
+    dicts for these fields, we need to convert them to the proper config classes
+    so that default values are populated correctly.
+
+    Without this conversion, dicts get converted to argparse.Namespace objects
+    which lack the default field values, causing AttributeError when vLLM code
+    tries to access those fields.
+    """
+    fields_by_name = {f.name: f for f in dataclasses.fields(AsyncEngineArgs)}
+
+    for key, value in list(merged.items()):
+        if not isinstance(value, dict) or key not in fields_by_name:
+            continue
+
+        hint = fields_by_name[key].type
+        if isinstance(hint, str):
+            continue
+
+        # Handle Optional[X] (Union[X, None]) -> X
+        origin = typing.get_origin(hint)
+        if origin is Union:
+            args = typing.get_args(hint)
+            hint = next((a for a in args if a is not type(None)), hint)
+
+        # Convert dict to dataclass if the field expects a dataclass type
+        if isinstance(hint, type) and dataclasses.is_dataclass(hint):
+            try:
+                merged[key] = hint(**value)
+            except Exception as e:
+                logger.warning(
+                    f"Failed to convert {key} dict to {hint.__name__}: {e}. "
+                    "Using dict as-is."
+                )
+
+    return merged
+
+
+def _dict_to_namespace(obj: Any) -> Any:
+    """Recursively converts dictionaries to argparse.Namespace."""
+    if isinstance(obj, dict):
+        return argparse.Namespace(**{k: _dict_to_namespace(v) for k, v in obj.items()})
+    elif isinstance(obj, list):
+        return [_dict_to_namespace(item) for item in obj]
+    else:
+        return obj
+
+
+def _get_vllm_engine_config(
+    llm_config: LLMConfig,
+) -> Tuple["AsyncEngineArgs", "VllmConfig"]:
+    engine_config = llm_config.get_engine_config()
+
+    # Resolve to local cache path if model was downloaded from S3/GCS mirror
+    # Only do this if mirror_config was specified (intentional S3/GCS download)
+    if engine_config.mirror_config:
+        from ray.llm._internal.common.utils.download_utils import (
+            get_model_location_on_disk,
+        )
+
+        local_path = get_model_location_on_disk(engine_config.actual_hf_model_id)
+        if local_path and local_path != engine_config.actual_hf_model_id:
+            engine_config.hf_model_id = local_path
+            logger.info(f"Resolved model from mirror to local path: {local_path}")
+
+    from vllm.usage.usage_lib import UsageContext
+
+    try:
+        async_engine_args = vllm.engine.arg_utils.AsyncEngineArgs(
+            **engine_config.get_initialization_kwargs()
+        )
+        vllm_engine_config = async_engine_args.create_engine_config(
+            usage_context=UsageContext.OPENAI_API_SERVER
+        )
+    except Exception as e:
+        # vLLM's ModelConfig is a pydantic dataclass; its ValidationError holds an
+        # unpicklable ArgsKwargs and cannot cross the Ray task boundary. Re-raise as
+        # a plain error so the real message propagates instead of a pickling failure.
+        raise RuntimeError(f"Failed to create vLLM engine config: {e}") from None
+    return async_engine_args, vllm_engine_config
+
+
+def _clear_current_platform_cache():
+    """Clear the cache of the current platform.
+
+    vllm current has an lru cache for getting device compatibility
+    that will not have the correct returned value if
+    CUDA_VISIBLE_DEVICES is not set properly. In RayLLM eventually
+    when we want to create the engine the env will be set properly,
+    but till then, upon the import of vllm somewhere
+    (which is a mystery) the lru cache will have the wrong value.
+    This function will clear the cache so that the next time the
+    cache is accessed, it will be re-evaluated.
+
+    Related issues:
+    https://github.com/vllm-project/vllm/issues/8402
+    https://github.com/vllm-project/vllm/issues/7890
+    """
+    from vllm.platforms import current_platform
+
+    # This check is just to future proof this implementation
+    # in case vllm removes their lru_cache decorator
+    if hasattr(current_platform.get_device_capability, "cache_clear"):
+        logger.info("Clearing the current platform cache ...")
+        current_platform.get_device_capability.cache_clear()
+
+
+class VLLMSleepConfig(BaseModel):
+    """vLLM-specific configuration for sleep operation."""
+
+    level: int = 1
+    """Sleep level:
+    - Level 1: Offload weights to CPU RAM, discard KV cache
+    - Level 2: Discard both model weights and KV cache (deeper sleep)
+    """
+
+    @field_validator("level")
+    @classmethod
+    def validate_level(cls, v: Any) -> int:
+        if v not in (1, 2):
+            raise ValueError("level must be 1 or 2")
+        return v
+
+
+class VLLMWakeupConfig(BaseModel):
+    """vLLM-specific configuration for wakeup operation."""
+
+    tags: Optional[List[str]] = None
+    """Optional tags to selectively wake up components:
+    - "weights": Restore model weights only
+    - "kv_cache": Restore KV cache only
+    - None: Restore everything
+    """
+
+    @field_validator("tags")
+    @classmethod
+    def validate_tags(cls, v: Any) -> Optional[List[str]]:
+        if v is not None:
+            valid_tags = {"weights", "kv_cache"}
+            for tag in v:
+                if tag not in valid_tags:
+                    raise ValueError(
+                        f"Invalid tag '{tag}'. Must be one of: {valid_tags}"
+                    )
+        return v
+
+
+class VLLMPauseConfig(BaseModel):
+    """vLLM-specific configuration for pause operation."""
+
+    mode: Literal["abort", "wait", "keep"] = "abort"
+    """Pause mode:
+    - "abort" (default): Abort all in-flight requests immediately.
+    - "wait": Wait for in-flight requests to complete before pausing.
+    - "keep": Freeze requests in queue; they resume on resume_generation().
+    """
+
+    clear_cache: bool = True
+    """Whether to clear KV and prefix caches after draining.
+    Set to False to preserve cache for faster resume.
+    """
+
+
+class VLLMEngine(LLMEngine):
+    def __init__(
+        self,
+        llm_config: LLMConfig,
+    ):
+        """Create a vLLM Engine class
+
+        Args:
+            llm_config: The llm configuration for this engine
+        """
+        super().__init__(llm_config)
+
+        self.llm_config = llm_config
+
+        if vllm is None:
+            raise ImportError(
+                "vLLM is not installed. Please install it with `pip install ray[llm]`."
+            )
+        assign_replica_kv_events_endpoint(self.llm_config)
+        self.llm_config.setup_engine_backend()
+
+        self._running = False
+        # Routing stats advertised to Serve's request router; populated in
+        # start() once the engine's KV-events endpoint is bound.
+        self._routing_stats: Dict[str, Any] = {}
+
+        # vLLM Integration points. Will be set through .start()
+        self._engine_client = None
+        self._oai_models: Optional["OpenAIServingModels"] = None
+        self._oai_serving_chat: Optional["OpenAIServingChat"] = None
+        self._oai_serving_completion: Optional["OpenAIServingCompletion"] = None
+        self._oai_serving_embedding: Optional["ServingEmbedding"] = None
+        self._oai_serving_transcription: Optional["OpenAIServingTranscription"] = None
+        self._oai_serving_scores: Optional["ServingScores"] = None
+        self._oai_serving_tokenization: Optional["OpenAIServingTokenization"] = None
+
+    async def build_asgi_app(self):
+        from vllm.entrypoints.openai.api_server import build_app, init_app_state
+
+        supported_tasks = ("generate",)
+        if hasattr(self._engine_client, "get_supported_tasks"):
+            supported_tasks = await self._engine_client.get_supported_tasks()
+
+        # Pass model_config so vLLM mounts the pooling routers (/pooling, /classify,
+        # /embed, /score) on the native ASGI app to enable direct streaming for pooling
+        # classify, embed, and score.
+        app = build_app(
+            self._vllm_args,
+            supported_tasks=supported_tasks,
+            model_config=self._engine_client.model_config,
+        )
+        await init_app_state(
+            self._engine_client,
+            app.state,
+            self._vllm_args,
+            supported_tasks=supported_tasks,
+        )
+        return app
+
+    async def start(self) -> None:
+        """Start the vLLM engine.
+
+        If the engine is already running, do nothing.
+        """
+
+        if self._running:
+            # The engine is already running!
+            logger.info("Skipping engine restart because the engine is already running")
+            return
+
+        from vllm.entrypoints.openai.api_server import init_app_state
+
+        callback = self.llm_config.get_or_create_callback()
+        await callback.run_callback("on_before_node_init")
+        if callback.ctx.run_init_node:
+            await initialize_node(self.llm_config)
+        await callback.run_callback("on_after_node_init")
+
+        (
+            vllm_engine_args,
+            vllm_frontend_args,
+            vllm_engine_config,
+        ) = self._prepare_engine_config(callback.ctx)
+
+        # Apply checkpoint info to the llm_config.
+        # This is needed for capturing model capabilities
+        # (e.g. supports vision, etc.) on the llm_config.
+        config = self.llm_config.get_engine_config()
+        self.llm_config.apply_checkpoint_info(
+            vllm_engine_config.model_config.model,
+            trust_remote_code=config.trust_remote_code,
+        )
+
+        self._engine_client = self._start_async_llm_engine(
+            vllm_engine_args,
+            vllm_engine_config,
+            callback.ctx.placement_group,
+        )
+
+        state = State()
+        # TODO (Kourosh): There might be some variables that needs protection?
+        merged = vllm_frontend_args.__dict__ | vllm_engine_args.__dict__
+
+        # Convert dict values to proper vLLM config classes (e.g., StructuredOutputsConfig)
+        # so that default field values are populated correctly.
+        merged = _convert_config_dicts(merged)
+
+        args = _dict_to_namespace(merged)
+        self._vllm_args = args
+
+        # Query supported tasks from the engine so init_app_state initializes the correct serving objects.
+        # Without this, vLLM falls back to 'generate' only.
+        init_kwargs: dict[str, Any] = dict(
+            state=state,
+            args=args,
+        )
+        if "supported_tasks" in inspect.signature(init_app_state).parameters:
+            if hasattr(self._engine_client, "get_supported_tasks"):
+                supported_tasks = await self._engine_client.get_supported_tasks()
+                init_kwargs["supported_tasks"] = supported_tasks
+
+        if "vllm_config" in inspect.signature(init_app_state).parameters:
+            init_kwargs["vllm_config"] = vllm_engine_config
+
+        await init_app_state(self._engine_client, **init_kwargs)
+
+        self._oai_models = getattr(state, "openai_serving_models", None)
+        self._oai_serving_chat = getattr(state, "openai_serving_chat", None)
+        self._oai_serving_completion = getattr(state, "openai_serving_completion", None)
+        self._oai_serving_embedding = getattr(state, "serving_embedding", None)
+        self._oai_serving_transcription = getattr(
+            state, "openai_serving_transcription", None
+        )
+        self._oai_serving_scores = getattr(state, "serving_scores", None)
+        self._oai_serving_tokenization = getattr(
+            state, "openai_serving_tokenization", None
+        )
+
+        self._validate_openai_serving_models()
+        self._validate_engine_client()
+
+        self._routing_stats = get_kv_event_routing_stats(
+            self.llm_config,
+            vllm_engine_config.cache_config.block_size,
+            vllm_engine_config.scheduler_config.max_num_batched_tokens,
+        )
+
+        self._running = True
+
+        logger.info("Started vLLM engine.")
+
+    def routing_stats(self) -> Dict[str, Any]:
+        """Returns KV event and replay endpoints for KV-aware routing."""
+        return self._routing_stats
+
+    def _validate_openai_serving_models(self):
+        assert self._oai_models is not None, "oai_models is not initialized"
+        assert hasattr(
+            self._oai_models, "lora_requests"
+        ), "oai_models must have a lora_requests attribute"
+        assert hasattr(
+            self._oai_models, "load_lora_adapter"
+        ), "oai_models must have a load_lora_adapter attribute"
+
+    @staticmethod
+    def _make_error(message: str) -> ErrorResponse:
+        return ErrorResponse(
+            error=ErrorInfo(message=message, type="invalid_request_error", code=400)
+        )
+
+    def _validate_openai_serving_chat(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_chat is None:
+            return self._make_error(
+                "This model does not support the 'generate' task. "
+                "The chat completion endpoint is not available for this model."
+            )
+
+    def _validate_openai_serving_completion(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_completion is None:
+            return self._make_error(
+                "This model does not support the 'generate' task. "
+                "The completion endpoint is not available for this model."
+            )
+
+    def _validate_openai_serving_embedding(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_embedding is None:
+            return self._make_error(
+                "This model does not support the 'embed' task. "
+                "The embedding endpoint is not available for this model."
+            )
+
+    def _validate_openai_serving_transcription(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_transcription is None:
+            return self._make_error(
+                "This model does not support the 'transcription' task. "
+                "The transcription endpoint is not available for this model."
+            )
+
+    def _validate_openai_serving_scores(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_scores is None:
+            return self._make_error(
+                "This model does not support the 'score' task. "
+                "The score endpoint is not available for this model."
+            )
+
+    def _validate_openai_serving_tokenization(self) -> Optional[ErrorResponse]:
+        if self._oai_serving_tokenization is None:
+            return self._make_error(
+                "This model does not support the 'tokenization' task. "
+                "The tokenization endpoint is not available for this model."
+            )
+
+    def _validate_engine_client(self):
+        assert hasattr(
+            self._engine_client, "check_health"
+        ), "engine_client must have a check_health attribute"
+
+    def _prepare_engine_config(
+        self, callback_ctx: CallbackCtx
+    ) -> Tuple["AsyncEngineArgs", "FrontendArgs", "VllmConfig"]:
+        """Prepare the engine config to start the engine.
+
+        Args:
+            callback_ctx: The callback context.
+
+        Returns:
+            A tuple of:
+                engine_args: The vLLM's internal engine arguments that is flattened.
+                frontend_args: The vLLM's internal frontend arguments that is flattened.
+                engine_config: The vLLM's internal engine config that is nested.
+        """
+
+        engine_config: VLLMEngineConfig = self.llm_config.get_engine_config()
+
+        # If the backend is anything other than CPU, we need to create the
+        # engine config on a task with hardware access.
+        if engine_config.accelerator.requires_remote_initialization:
+            accelerator = engine_config.accelerator
+            accelerator_type = self.llm_config.accelerator_type
+
+            # Initialize options required for the remote task and hardware backend
+            remote_options = {
+                "num_cpus": 0,
+                "runtime_env": callback_ctx.runtime_env,
+                "scheduling_strategy": PlacementGroupSchedulingStrategy(
+                    placement_group=callback_ctx.placement_group,
+                ),
+                **accelerator.get_remote_options(accelerator_type),
+            }
+
+            ref = (
+                ray.remote(_get_vllm_engine_config)
+                .options(**remote_options)
+                .remote(self.llm_config)
+            )
+            vllm_engine_args, vllm_engine_config = ray.get(ref)
+        else:
+            vllm_engine_args, vllm_engine_config = _get_vllm_engine_config(
+                self.llm_config
+            )
+
+        vllm_frontend_args = FrontendArgs(**engine_config.frontend_kwargs)
+        return vllm_engine_args, vllm_frontend_args, vllm_engine_config
+
+    def _start_async_llm_engine(
+        self,
+        vllm_engine_args: "AsyncEngineArgs",
+        vllm_engine_config: "VllmConfig",
+        placement_group: PlacementGroup,
+    ) -> "EngineClient":
+        """Creates an async LLM engine from the engine arguments."""
+
+        from vllm.v1.engine.async_llm import AsyncLLM
+        from vllm.v1.executor.abstract import Executor
+
+        vllm_engine_config.parallel_config.placement_group = placement_group
+
+        _clear_current_platform_cache()
+
+        custom_stat_loggers = None
+        if self.llm_config.log_engine_metrics:
+            from vllm.v1.metrics.ray_wrappers import RayPrometheusStatLogger
+
+            # V1 AsyncLLM does not yet support add_logger: https://github.com/vllm-project/vllm/issues/17702
+            # Use `disable_log_stats: False` and `log_engine_metrics: False` as
+            # a workaround to enable PrometheusStatLogger instead.
+            custom_stat_loggers = [RayPrometheusStatLogger]
+
+        executor_class = Executor.get_class(vllm_engine_config)
+        logger.info(f"Using executor class: {executor_class}")
+        # Report per-request token progress to the deployment's KV router actor,
+        # but only on KV-aware deployments: elsewhere the actor never exists and
+        # resolving it per request would block the engine's event loop.
+        engine_cls = AsyncLLM
+        if is_kv_aware(self.llm_config):
+            engine_cls = enable_token_tracking(AsyncLLM)
+        engine_client = engine_cls(
+            vllm_config=vllm_engine_config,
+            executor_class=executor_class,
+            log_requests=vllm_engine_args.enable_log_requests,
+            log_stats=not vllm_engine_args.disable_log_stats,
+            stat_loggers=custom_stat_loggers,
+        )
+
+        return engine_client
+
+    async def resolve_lora(self, disk_lora_model: DiskMultiplexConfig):
+        from vllm.entrypoints.serve.lora.protocol import LoadLoRAAdapterRequest
+
+        self._validate_openai_serving_models()
+
+        if disk_lora_model.model_id in self._oai_models.lora_requests:
+            # Lora is already loaded, return
+            return
+
+        lora_request = await self._oai_models.load_lora_adapter(  # type: ignore[attr-defined]
+            request=LoadLoRAAdapterRequest(
+                lora_name=disk_lora_model.model_id,
+                lora_path=disk_lora_model.local_path,
+            )
+        )
+
+        if isinstance(lora_request, VLLMErrorResponse):
+            raise ValueError(f"Failed to load lora model: {lora_request.error.message}")
+
+    @staticmethod
+    def _make_error_response(
+        serving: Any,
+        exc: Exception,
+    ) -> ErrorResponse:
+        """Convert an exception to an ErrorResponse and map exception types to
+        the appropriate HTTP status codes (e.g. VLLMValidationError -> 400).
+        """
+        try:
+            vllm_error = serving.create_error_response(exc)
+            return ErrorResponse(error=ErrorInfo(**vllm_error.error.model_dump()))
+        except Exception:
+            raise exc  # re-raise the original so it surfaces as a 500
+
+    async def chat(
+        self,
+        request: ChatCompletionRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[Union[str, ChatCompletionResponse, ErrorResponse], None]:
+        if error := self._validate_openai_serving_chat():
+            yield error
+            return
+
+        raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
+        raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
+            raw_request_info
+        )
+        try:
+            chat_response = await self._oai_serving_chat.create_chat_completion(  # type: ignore[attr-defined]
+                request,
+                raw_request=raw_request,
+            )
+        except ValueError as e:
+            yield self._make_error_response(self._oai_serving_chat, e)
+            return
+
+        if isinstance(chat_response, AsyncGenerator):
+            async for response in chat_response:
+                if not isinstance(response, str):
+                    raise ValueError(
+                        f"Expected create_chat_completion to return a stream of strings, got an item with type {type(response)}"
+                    )
+                yield response
+        else:
+            if isinstance(chat_response, VLLMErrorResponse):
+                yield ErrorResponse(error=ErrorInfo(**chat_response.error.model_dump()))
+            else:
+                yield ChatCompletionResponse(**chat_response.model_dump())
+
+    async def completions(
+        self,
+        request: CompletionRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[Union[str, CompletionResponse, ErrorResponse], None]:
+        if error := self._validate_openai_serving_completion():
+            yield error
+            return
+
+        raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
+        raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
+            raw_request_info
+        )
+        try:
+            completion_response = await self._oai_serving_completion.create_completion(  # type: ignore[attr-defined]
+                request,
+                raw_request=raw_request,
+            )
+        except ValueError as e:
+            yield self._make_error_response(self._oai_serving_completion, e)
+            return
+
+        if isinstance(completion_response, AsyncGenerator):
+            async for response in completion_response:
+                if not isinstance(response, str):
+                    raise ValueError(
+                        f"Expected create_completion to return a stream of strings, got an item with type {type(response)}"
+                    )
+                yield response
+        else:
+            if isinstance(completion_response, VLLMErrorResponse):
+                yield ErrorResponse(
+                    error=ErrorInfo(**completion_response.error.model_dump())
+                )
+            else:
+                yield CompletionResponse(**completion_response.model_dump())
+
+    async def embeddings(
+        self,
+        request: EmbeddingRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[Union[EmbeddingResponse, ErrorResponse], None]:
+        if error := self._validate_openai_serving_embedding():
+            yield error
+            return
+
+        raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
+        raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
+            raw_request_info
+        )
+        try:
+            embedding_response = await self._oai_serving_embedding(
+                request,
+                raw_request=raw_request,
+            )
+        except ValueError as e:
+            yield self._make_error_response(self._oai_serving_embedding, e)
+            return
+
+        # vLLM 0.18+ returns a starlette Response object
+        content = json.loads(embedding_response.body)
+        yield EmbeddingResponse(**content)
+
+    async def transcriptions(
+        self,
+        request: TranscriptionRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[Union[str, TranscriptionResponse, ErrorResponse], None]:
+        if error := self._validate_openai_serving_transcription():
+            yield error
+            return
+
+        # Extract audio data from the request file
+        audio_data = await request.file.read()
+
+        raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
+        raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
+            raw_request_info
+        )
+        try:
+            transcription_response = await self._oai_serving_transcription.create_transcription(  # type: ignore[attr-defined]
+                audio_data,
+                request,
+                raw_request=raw_request,
+            )
+        except ValueError as e:
+            yield self._make_error_response(self._oai_serving_transcription, e)
+            return
+
+        if isinstance(transcription_response, AsyncGenerator):
+            async for response in transcription_response:
+                if not isinstance(response, str):
+                    raise ValueError(
+                        f"Expected create_transcription to return a stream of strings, got an item with type {type(response)}"
+                    )
+                yield response
+        else:
+            if isinstance(transcription_response, VLLMErrorResponse):
+                yield ErrorResponse(
+                    error=ErrorInfo(**transcription_response.error.model_dump())
+                )
+            else:
+                yield TranscriptionResponse(**transcription_response.model_dump())
+
+    async def score(
+        self,
+        request: ScoreRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[Union[ScoreResponse, ErrorResponse], None]:
+        if error := self._validate_openai_serving_scores():
+            yield error
+            return
+
+        raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
+        raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
+            raw_request_info
+        )
+        try:
+            assert self._oai_serving_scores is not None
+            score_response = await self._oai_serving_scores(
+                request,
+                raw_request=raw_request,
+            )
+        except ValueError as e:
+            yield self._make_error_response(self._oai_serving_scores, e)
+            return
+
+        content = json.loads(score_response.body)
+        yield ScoreResponse(**content)
+
+    async def tokenize(
+        self,
+        request: TokenizeRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[Union[TokenizeResponse, ErrorResponse], None]:
+        if error := self._validate_openai_serving_tokenization():
+            yield error
+            return
+
+        raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
+        raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
+            raw_request_info
+        )
+        try:
+            tokenize_response = await self._oai_serving_tokenization.create_tokenize(
+                request,
+                raw_request=raw_request,
+            )
+        except ValueError as e:
+            yield self._make_error_response(self._oai_serving_tokenization, e)
+            return
+
+        if isinstance(tokenize_response, VLLMErrorResponse):
+            yield ErrorResponse(error=ErrorInfo(**tokenize_response.error.model_dump()))
+        else:
+            yield TokenizeResponse(**tokenize_response.model_dump())
+
+    async def detokenize(
+        self,
+        request: DetokenizeRequest,
+        raw_request_info: Optional[RawRequestInfo] = None,
+    ) -> AsyncGenerator[Union[DetokenizeResponse, ErrorResponse], None]:
+        if error := self._validate_openai_serving_tokenization():
+            yield error
+            return
+
+        raw_request_info = _canonicalize_request_id_header(request, raw_request_info)
+        raw_request: Optional[Request] = RawRequestInfo.to_starlette_request_optional(
+            raw_request_info
+        )
+        try:
+            detokenize_response = (
+                await self._oai_serving_tokenization.create_detokenize(
+                    request,
+                    raw_request=raw_request,
+                )
+            )
+        except ValueError as e:
+            yield self._make_error_response(self._oai_serving_tokenization, e)
+            return
+
+        if isinstance(detokenize_response, VLLMErrorResponse):
+            yield ErrorResponse(
+                error=ErrorInfo(**detokenize_response.error.model_dump())
+            )
+        else:
+            yield DetokenizeResponse(**detokenize_response.model_dump())
+
+    async def check_health(self) -> None:
+        assert self._engine_client is not None, "engine_client is not initialized"
+
+        try:
+            await self._engine_client.check_health()
+        except BaseException as e:
+            logger.error("Healthcheck failed. The replica will be restarted")
+            raise e from None
+
+    async def reset_prefix_cache(self) -> None:
+        assert self._engine_client is not None, "engine_client is not initialized"
+        await self._engine_client.reset_prefix_cache()
+
+    async def sleep(self, **kwargs: Any) -> None:
+        """Put the vLLM engine to sleep.
+
+        Args:
+            **kwargs: Options parsed into VLLMSleepConfig.
+                - level (int): Sleep level (1 or 2). Default 1.
+        """
+        assert self._engine_client is not None, "engine_client is not initialized"
+        config = VLLMSleepConfig(**kwargs)
+        await self._engine_client.sleep(level=config.level)
+
+    async def wakeup(self, **kwargs: Any) -> None:
+        """Wake up the vLLM engine from sleep mode.
+
+        Args:
+            **kwargs: Options parsed into VLLMWakeupConfig.
+                - tags (List[str], optional): Components to wake up.
+        """
+        assert self._engine_client is not None, "engine_client is not initialized"
+        config = VLLMWakeupConfig(**kwargs)
+        await self._engine_client.wake_up(tags=config.tags)
+
+    async def is_sleeping(self) -> bool:
+        """Check whether the vLLM engine is currently sleeping.
+
+        Returns:
+            True if the engine is sleeping, False otherwise.
+        """
+        assert self._engine_client is not None, "engine_client is not initialized"
+        return await self._engine_client.is_sleeping()
+
+    async def pause(self, **kwargs: Any) -> None:
+        """Pause generation on the vLLM engine.
+
+        This halts generation/encoding requests while keeping model weights
+        in GPU memory. New requests are blocked until resume is called.
+
+        Args:
+            **kwargs: Options parsed into VLLMPauseConfig.
+                - mode (str): "abort" (default), "wait", or "keep".
+                - clear_cache (bool): Clear KV cache after draining. Default True.
+        """
+        assert self._engine_client is not None, "engine_client is not initialized"
+        config = VLLMPauseConfig(**kwargs)
+        await self._engine_client.pause_generation(
+            mode=config.mode,
+            clear_cache=config.clear_cache,
+        )
+
+    async def resume(self, **kwargs: Any) -> None:
+        """Resume generation on the vLLM engine after pause.
+
+        Args:
+            **kwargs: Reserved for future options.
+        """
+        assert self._engine_client is not None, "engine_client is not initialized"
+        await self._engine_client.resume_generation()
+
+    async def is_paused(self) -> bool:
+        """Check whether the vLLM engine is currently paused.
+
+        Returns:
+            True if the engine is paused, False otherwise.
+        """
+        assert self._engine_client is not None, "engine_client is not initialized"
+        return await self._engine_client.is_paused()
+
+    async def start_profile(self) -> None:
+        assert self._engine_client is not None, "engine_client is not initialized"
+        await self._engine_client.start_profile()
+
+    async def stop_profile(self) -> None:
+        assert self._engine_client is not None, "engine_client is not initialized"
+        await self._engine_client.stop_profile()
+
+    async def collective_rpc(
+        self,
+        method: str,
+        timeout: Optional[float] = None,
+        args: tuple = (),
+        kwargs: Optional[dict] = None,
+    ) -> list:
+        """Execute a collective RPC call on all vLLM workers.
+
+        This is used for RLHF workflows where a trainer needs to execute
+        methods on all TP/PP workers (e.g., for weight synchronization).
+
+        Args:
+            method: Name of the worker method to execute.
+            timeout: Maximum time in seconds to wait for execution.
+            args: Positional arguments to pass to the worker method.
+            kwargs: Keyword arguments to pass to the worker method.
+
+        Returns:
+            A list containing the results from each worker.
+        """
+        assert self._engine_client is not None, "engine_client is not initialized"
+        return await self._engine_client.collective_rpc(
+            method=method,
+            timeout=timeout,
+            args=args,
+            kwargs=kwargs or {},
+        )

python/ray/llm/_internal/serve/engines/vllm/vllm_models.py

@@ -0,0 +1,361 @@
+import dataclasses
+import os
+from typing import Any, Dict, List, Optional
+
+from pydantic import ConfigDict, Field, PrivateAttr, field_validator, model_validator
+from vllm.engine.arg_utils import AsyncEngineArgs
+from vllm.entrypoints.openai.cli_args import FrontendArgs
+
+from ray.llm._internal.common.base_pydantic import BaseModelExtended
+from ray.llm._internal.common.placement import PlacementGroupConfig
+from ray.llm._internal.common.utils.cloud_utils import CloudMirrorConfig, is_remote_path
+from ray.llm._internal.common.utils.import_utils import try_import
+from ray.llm._internal.serve.constants import (
+    ALLOW_NEW_PLACEMENT_GROUPS_IN_DEPLOYMENT,
+    ENV_VARS_TO_PROPAGATE,
+)
+from ray.llm._internal.serve.core.configs.accelerators import (
+    AcceleratorBackend,
+    AnyAcceleratorConfig,
+    CPUAccelerator,
+    CPUConfig,
+    GPUAccelerator,
+    TPUAccelerator,
+    TPUConfig,
+    format_ray_accelerator_resource,
+)
+from ray.llm._internal.serve.core.configs.llm_config import (
+    AcceleratorType,
+    LLMConfig,
+)
+from ray.llm._internal.serve.observability.logging import get_logger
+from ray.util.placement_group import (
+    PlacementGroup,
+    get_current_placement_group,
+    placement_group_table,
+)
+
+# The key for the kv_transfer_params in the internal metadata.
+KV_TRANSFER_PARAMS_KEY = "kv_transfer_params"
+vllm = try_import("vllm")
+logger = get_logger(__name__)
+
+# Executor backend constants
+EXECUTOR_BACKEND_RAY = "ray"
+EXECUTOR_BACKEND_MP = "mp"
+
+
+class VLLMEngineConfig(BaseModelExtended):
+    model_config = ConfigDict(
+        use_enum_values=True,
+    )
+
+    model_id: str = Field(
+        description="The identifier for the model. This is the id that will be used to query the model.",
+    )
+    hf_model_id: Optional[str] = Field(
+        None, description="The Hugging Face model identifier."
+    )
+    mirror_config: Optional[CloudMirrorConfig] = Field(
+        None,
+        description="Configuration for cloud storage mirror. This is for where the weights are downloaded from.",
+    )
+    accelerator_type: Optional[AcceleratorType] = Field(
+        None,
+        description="The type of accelerator to use. This is used to determine the placement group strategy.",
+    )
+    accelerator_config: Optional[AnyAcceleratorConfig] = Field(
+        default=None,
+        description=(
+            "Hardware-specific configuration parameters for the chosen accelerator. "
+            "The expected schema is dynamically typed based on the 'kind' discriminator."
+        ),
+    )
+
+    placement_group_config: Optional[Dict[str, Any]] = Field(
+        default=None,
+        description=(
+            "Ray placement group configuration for scheduling vLLM engine workers. "
+            "Defines resource bundles and placement strategy for multi-node deployments. "
+            "Defaults to PACK strategy with automatic bundle generation based on TP/PP sizes."
+        ),
+    )
+
+    @field_validator("placement_group_config")
+    @classmethod
+    def validate_placement_group_config(cls, value):
+        if value is None:
+            return None
+        # Validate through PlacementGroupConfig, then dump back to dict
+        validated = PlacementGroupConfig(**value)
+        return validated.model_dump(exclude_unset=True)
+
+    runtime_env: Optional[Dict[str, Any]] = None
+    engine_kwargs: Dict[str, Any] = {}
+    frontend_kwargs: Dict[str, Any] = {}
+
+    _accelerator: AcceleratorBackend = PrivateAttr(default=None)
+
+    @model_validator(mode="after")
+    def _build_accelerator(self):
+        """Instantiates the accelerator backend based on the resolved config."""
+        cfg = self.accelerator_config
+
+        if self.accelerator_type and isinstance(cfg, CPUConfig):
+            raise ValueError(
+                f"accelerator_type='{self.accelerator_type}' cannot be used with "
+                "CPU-only configurations. Either remove accelerator_type, or provide an accelerator_config."
+            )
+
+        # LLMConfig has already resolved and validated accelerator_config
+        if isinstance(cfg, TPUConfig):
+            self._accelerator = TPUAccelerator(cfg)
+        elif isinstance(cfg, CPUConfig):
+            self._accelerator = CPUAccelerator()
+        else:
+            # Default to GPU if it's GPUConfig or isn't set
+            self._accelerator = GPUAccelerator()
+        return self
+
+    @property
+    def accelerator(self) -> AcceleratorBackend:
+        return self._accelerator
+
+    @property
+    def actual_hf_model_id(self) -> str:
+        return self.hf_model_id or self.model_id
+
+    @property
+    def trust_remote_code(self) -> bool:
+        return self.engine_kwargs.get("trust_remote_code", False)
+
+    def get_initialization_kwargs(self) -> dict:
+        """
+        Get kwargs that will be actually passed to the LLMInitializer
+        constructor.
+        """
+        engine_kwargs = self.engine_kwargs.copy()
+
+        if "model" in engine_kwargs or "served_model_name" in engine_kwargs:
+            raise ValueError(
+                "model or served_model_name is not allowed in engine_kwargs when using Ray Serve LLM. Please use `model_loading_config` in LLMConfig instead."
+            )
+
+        engine_kwargs["model"] = self.actual_hf_model_id
+        engine_kwargs["served_model_name"] = [self.model_id]
+
+        # Handle distributed_executor_backend based on backend type
+        if isinstance(self.accelerator, CPUAccelerator):
+            executor_backend = EXECUTOR_BACKEND_MP
+        else:
+            executor_backend = EXECUTOR_BACKEND_RAY
+
+        if (
+            "distributed_executor_backend" in engine_kwargs
+            and engine_kwargs["distributed_executor_backend"] != executor_backend
+            and executor_backend == EXECUTOR_BACKEND_RAY
+        ):
+            raise ValueError(
+                "distributed_executor_backend != 'ray' is not allowed in engine_kwargs when using Ray Serve LLM Configs."
+            )
+
+        engine_kwargs["distributed_executor_backend"] = executor_backend
+
+        if "enable_log_requests" not in engine_kwargs:
+            engine_kwargs["enable_log_requests"] = False
+
+        return engine_kwargs
+
+    def get_runtime_env_with_local_env_vars(self) -> dict:
+        runtime_env = self.runtime_env or {}
+        runtime_env.setdefault("env_vars", {})
+        env_vars = runtime_env["env_vars"]
+
+        # Propagate env vars to the runtime env
+        for env_var in ENV_VARS_TO_PROPAGATE:
+            if env_var in os.environ:
+                env_vars[env_var] = os.getenv(env_var)
+
+        if "VLLM_RAY_PER_WORKER_GPUS" not in env_vars:
+            fractional_gpu = self._detect_fractional_gpu_from_pg(
+                self.placement_group_config
+            )
+            if fractional_gpu is not None:
+                env_vars["VLLM_RAY_PER_WORKER_GPUS"] = str(fractional_gpu)
+        return runtime_env
+
+    @classmethod
+    def from_llm_config(cls, llm_config: LLMConfig) -> "VLLMEngineConfig":
+        """Converts the LLMConfig to a VLLMEngineConfig."""
+        # Set up the model downloading configuration.
+        hf_model_id, mirror_config = None, None
+        if llm_config.model_loading_config.model_source is None:
+            hf_model_id = llm_config.model_id
+        elif isinstance(llm_config.model_loading_config.model_source, str):
+            model_source = llm_config.model_loading_config.model_source
+            if is_remote_path(model_source):
+                # Remote URIs (s3://, gs://, …) are download addresses,
+                # not HuggingFace IDs.  Using the URI verbatim as
+                # hf_model_id propagates the scheme and slashes into the
+                # cache directory name (``models--s3:----bucket--…``).
+                # Use the user-supplied model_id as the identifier and
+                # treat the URI as a bucket mirror instead.
+                hf_model_id = llm_config.model_id
+                mirror_config = CloudMirrorConfig(bucket_uri=model_source)
+            else:
+                hf_model_id = model_source
+        else:
+            # If it's a CloudMirrorConfig (or subtype)
+            mirror_config = llm_config.model_loading_config.model_source
+
+        all_engine_kwargs = llm_config.engine_kwargs.copy()
+        engine_kwargs = {}
+        frontend_kwargs = {}
+
+        # Get field names from dataclasses
+        frontend_field_names = {
+            field.name for field in dataclasses.fields(FrontendArgs)
+        }
+        async_engine_field_names = {
+            field.name for field in dataclasses.fields(AsyncEngineArgs)
+        }
+
+        for key, value in all_engine_kwargs.items():
+            if key in frontend_field_names:
+                frontend_kwargs[key] = value
+            elif key in async_engine_field_names:
+                engine_kwargs[key] = value
+            else:
+                raise ValueError(f"Unknown engine argument: {key}")
+
+        # placement_group_config is already validated and stored as dict in LLMConfig
+        placement_group_config = llm_config.placement_group_config
+        return VLLMEngineConfig(
+            model_id=llm_config.model_id,
+            hf_model_id=hf_model_id,
+            mirror_config=mirror_config,
+            accelerator_type=llm_config.accelerator_type,
+            accelerator_config=llm_config.accelerator_config,
+            engine_kwargs=engine_kwargs,
+            frontend_kwargs=frontend_kwargs,
+            runtime_env=llm_config.runtime_env,
+            placement_group_config=placement_group_config,
+        )
+
+    @property
+    def tensor_parallel_degree(self) -> int:
+        return self.engine_kwargs.get("tensor_parallel_size", 1)
+
+    @property
+    def pipeline_parallel_degree(self) -> int:
+        return self.engine_kwargs.get("pipeline_parallel_size", 1)
+
+    @property
+    def num_devices(self) -> int:
+        return self.tensor_parallel_degree * self.pipeline_parallel_degree
+
+    @property
+    def placement_strategy(self) -> str:
+        # Use custom strategy if placement_group_config is provided
+        if self.placement_group_config:
+            return self.placement_group_config.get("strategy", "PACK")
+        # Default to PACK (cross-node best-effort placement)
+        # DP deployments overridden to STRICT_PACK in Serve config
+        return "PACK"
+
+    @property
+    def placement_bundles(self) -> List[Dict[str, float]]:
+        if self.placement_group_config:
+            bundle_per_worker = self.placement_group_config.get("bundle_per_worker")
+
+            if bundle_per_worker is not None:
+                # Expand bundle_per_worker to num_devices bundles
+                bundles = []
+                for _ in range(self.num_devices):
+                    bundle = bundle_per_worker.copy()
+                    if self.accelerator_type:
+                        res_key = format_ray_accelerator_resource(self.accelerator_type)
+                        bundle.setdefault(res_key, 0.001)
+                    bundles.append(bundle)
+                return bundles
+
+            # Otherwise use explicit bundles list
+            bundles = []
+            explicit_bundles = self.placement_group_config.get("bundles") or []
+            for bundle_dict in explicit_bundles:
+                bundle = bundle_dict.copy()
+                if self.accelerator_type:
+                    # Use setdefault to add accelerator hint WITHOUT overriding explicit user values
+                    res_key = format_ray_accelerator_resource(self.accelerator_type)
+                    bundle.setdefault(res_key, 0.001)
+                bundles.append(bundle)
+            return bundles
+
+        # Default bundles based on the accelerator backend.
+        return self.accelerator.default_bundles(
+            num_devices=self.num_devices, accelerator_type_str=self.accelerator_type
+        )
+
+    def get_or_create_pg(self) -> PlacementGroup:
+        """Gets or creates a placement group.
+
+        If we are already in a placement group, return the existing placement group.
+        Else, delegate PG creation to the accelerator backend.
+        """
+        dp_rank = self.engine_kwargs.get("data_parallel_rank", None)
+        pg = get_current_placement_group()
+        if pg:
+            logger.debug(
+                "Using existing placement group %s, details: %s",
+                pg.id,
+                placement_group_table(pg),
+            )
+        else:
+            if not ALLOW_NEW_PLACEMENT_GROUPS_IN_DEPLOYMENT:
+                raise RuntimeError(
+                    "Creating new placement groups is not allowed. "
+                    "Change RAYLLM_ALLOW_NEW_PLACEMENT_GROUPS_IN_DEPLOYMENT "
+                    "if this is not intended."
+                )
+            name = "" if dp_rank is None else f"dp_{dp_rank}"
+
+            pg = self.accelerator.create_placement_group(
+                bundles=self.placement_bundles,
+                strategy=self.placement_strategy,
+                name=name,
+                accelerator_type_str=self.accelerator_type,
+            )
+
+            logger.info(f"Using new placement group {pg}. {placement_group_table(pg)}")
+
+        return pg
+
+    @staticmethod
+    def _detect_fractional_gpu_from_pg(
+        placement_group_config: Optional[Dict[str, Any]]
+    ) -> Optional[float]:
+        if not placement_group_config:
+            return None
+
+        # Check bundle_per_worker first
+        bundle_per_worker = placement_group_config.get("bundle_per_worker")
+        if bundle_per_worker:
+            gpu_value = bundle_per_worker.get("GPU", 0)
+            if 0 < gpu_value < 1:
+                return gpu_value
+            return None
+
+        # Fall back to bundles list
+        bundles = placement_group_config.get("bundles") or []
+
+        for bundle in bundles:
+            if "GPU" not in bundle:
+                continue
+
+            gpu_value = bundle["GPU"]
+            if gpu_value <= 0 or gpu_value >= 1:
+                return None
+
+            return gpu_value
+
+        return None