[serve][llm] Tokenize pre-routing requests in-process via the engine's vLLM renderer

Replace the per-request /tokenize RPC with an in-process tokenizer that
replays vLLM's ServingTokenization.create_tokenize on the engine's own
OnlineRenderer, so routing token ids are byte-identical to the replica's
prefill tokens without contending with inference. vLLM imports are hoisted
to module top (needs OnlineRenderer, vLLM >= 0.25) and the ingress router
imports the module lazily so non-KV paths never pull the renderer in.

Signed-off-by: Jeffrey Wang <jeffreywang@anyscale.com>
Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>

Author: jeffreywang88

python/ray/llm/_internal/serve/core/ingress/builder.py

@@ -104,7 +104,7 @@ def _build_openai_ingress_request_router(
     )
     return deployment.bind(
         server=server,
-        pre_routing_tokenization=is_kv_aware(llm_config),
+        llm_config=llm_config if is_kv_aware(llm_config) else None,
     )
 
 

python/ray/llm/_internal/serve/core/ingress/router.py

@@ -1,20 +1,24 @@
+import asyncio
 import json
 from types import SimpleNamespace
-from typing import List, Optional, Tuple
+from typing import TYPE_CHECKING, List, Optional, Tuple
 
 from fastapi import FastAPI, HTTPException, Request
 
 from ray import serve
 from ray.llm._internal.serve.observability.logging import get_logger
-from ray.llm._internal.serve.routing_policies.kv_aware.tokenizer import (
+from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
     REQUEST_TOKEN_IDS_KWARG,
-    TokenizeError,
-    Tokenizer,
 )
 from ray.serve._private.http_util import _matches_session_id_header
 from ray.serve.exceptions import DeploymentUnavailableError
 from ray.serve.handle import DeploymentHandle
 
+# Type-only import as LLMConfig transitively pulls in vLLM. This file should
+# remain engine-agnostic.
+if TYPE_CHECKING:
+    from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+
 logger = get_logger(__name__)
 
 _BODY_TRUNCATED_HEADER = "x-body-truncated"
@@ -105,13 +109,23 @@ class LLMRouter:
     _warned_no_routing_key: bool = False
 
     async def __init__(
-        self, server: DeploymentHandle, pre_routing_tokenization: bool = False
+        self,
+        server: DeploymentHandle,
+        llm_config: Optional["LLMConfig"] = None,
     ):
         self._handle: DeploymentHandle = server
         self._handle._init()
-        # Pre-routing tokenization is only useful to a KV-aware request router,
-        # which scores replicas based on the prompt token IDs.
-        self._tokenizer = Tokenizer(self._handle) if pre_routing_tokenization else None
+        self._tokenizer = None
+        # A non-None llm_config signals pre-routing tokenization, which the
+        # builder binds only for a KV-aware request router.
+        if llm_config is not None:
+            # Lazy import: this module pulls in vLLM's renderer;
+            # keep it off the non-KV ingress import path.
+            from ray.llm._internal.serve.routing_policies.kv_aware.tokenizer import (
+                Tokenizer,
+            )
+
+            self._tokenizer = await asyncio.to_thread(Tokenizer, llm_config)
 
     @router_app.post("/internal/route")
     async def route(self, request: Request):
@@ -133,6 +147,10 @@ async def route(self, request: Request):
         # body has no routing payload, so fall back to token-less routing.
         request_token_ids = None
         if self._tokenizer is not None and routing_payload is not None:
+            from ray.llm._internal.serve.routing_policies.kv_aware.tokenizer import (
+                TokenizeError,
+            )
+
             try:
                 request_token_ids = await self._tokenizer.tokenize(
                     vars(routing_payload)

python/ray/llm/_internal/serve/routing_policies/kv_aware/constants.py

@@ -1,3 +1,6 @@
+# choose_replica kwarg carrying the prompt token IDs to KV-aware routers.
+REQUEST_TOKEN_IDS_KWARG = "request_token_ids"
+
 # experimental_configs key overriding the per-node base port.
 KV_EVENTS_PORT_BASE_KEY = "KV_EVENTS_PORT_BASE"
 DEFAULT_KV_EVENTS_PORT_BASE = 5557

python/ray/llm/_internal/serve/routing_policies/kv_aware/kv_aware_router.py

@@ -5,13 +5,13 @@
 import ray
 from ray.actor import ActorHandle
 from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
+from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
+    REQUEST_TOKEN_IDS_KWARG,
+)
 from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
     KV_ROUTER_ACTOR_NAME,
     get_worker_id,
 )
-from ray.llm._internal.serve.routing_policies.kv_aware.tokenizer import (
-    REQUEST_TOKEN_IDS_KWARG,
-)
 from ray.serve._private.constants import (
     SERVE_DEPLOYMENT_ACTOR_PREFIX,
     SERVE_LOGGER_NAME,

python/ray/llm/_internal/serve/routing_policies/kv_aware/tokenizer.py

@@ -1,23 +1,30 @@
 from typing import Any, Dict, List, Optional, Union
 
+import jinja2
+from vllm.entrypoints.chat_utils import load_chat_template
+from vllm.entrypoints.openai.cli_args import FrontendArgs
+from vllm.renderers import renderer_from_config
+from vllm.renderers.inputs.preprocess import extract_prompt_components
+from vllm.renderers.online_renderer import OnlineRenderer
+
+from ray.llm._internal.serve.core.configs.llm_config import LLMConfig
 from ray.llm._internal.serve.core.configs.openai_api_models import (
-    ErrorResponse,
     TokenizeChatRequest,
     TokenizeCompletionRequest,
 )
+from ray.llm._internal.serve.engines.vllm.vllm_engine import (
+    _get_vllm_engine_config,
+)
 from ray.llm._internal.serve.observability.logging import get_logger
-from ray.serve.handle import DeploymentHandle
 
 logger = get_logger(__name__)
 
-# choose_replica kwarg carrying the prompt token IDs to KV-aware routers.
-REQUEST_TOKEN_IDS_KWARG = "request_token_ids"
-
 
 class TokenizeError(Exception):
-    """The ``/tokenize`` endpoint rejected the request.
+    """The request was rejected the same way vLLM's native ASGI route
+    ``/tokenize`` would reject it.
 
-    Carries vLLM's HTTP ``status_code``, ``message`` and error ``type``.
+    Carries the HTTP ``status_code``, ``message`` and error ``type``.
     """
 
     def __init__(self, message: str, *, status_code: int, type: str):
@@ -27,15 +34,77 @@ def __init__(self, message: str, *, status_code: int, type: str):
         self.type = type
 
 
+def build_tokenize_request(
+    payload: Dict[str, Any]
+) -> Optional[Union[TokenizeChatRequest, TokenizeCompletionRequest]]:
+    """Build the Tokenize* request for ``payload``, forwarding only the fields
+    the engine renders the prompt from so routing ids match the prefill tokens.
+
+    Returns ``None`` (caller falls back to token-less routing) for a body with
+    no single string prompt, e.g. a batch ``prompt`` list, since KV-aware
+    routing scores one request on one token sequence.
+
+    TODO (jeffreywang): Support multi-prompt tokenization.
+    """
+    try:
+        if "messages" in payload:
+            return TokenizeChatRequest.model_validate(
+                {
+                    k: v
+                    for k, v in payload.items()
+                    if k in TokenizeChatRequest.model_fields
+                }
+            )
+        if "prompt" in payload:
+            if not isinstance(payload["prompt"], str):
+                return None
+            return TokenizeCompletionRequest.model_validate(
+                {
+                    k: v
+                    for k, v in payload.items()
+                    if k in TokenizeCompletionRequest.model_fields
+                }
+            )
+        # Unreachable: LLMRouter only routes bodies with messages or a prompt.
+        logger.warning(
+            "Tokenizer got a payload with neither messages nor prompt; "
+            "falling back to token-less routing."
+        )
+        return None
+    except Exception as e:
+        logger.warning("Unsupported tokenize request, falling back: %s", e)
+        return None
+
+
 class Tokenizer:
-    """Tokenizes incoming requests via the replica's ``/tokenize`` endpoint.
+    """Tokenizes requests with vLLM's ``OnlineRenderer``.
+
+    Configured from the deployment's frontend args so the tokenizer, chat
+    template, and trust policy match the engine's.
 
     Args:
-        handle: A handle to the LLMServer deployment.
+        llm_config: The deployment's LLM config.
     """
 
-    def __init__(self, handle: DeploymentHandle):
-        self._handle = handle
+    def __init__(self, llm_config: LLMConfig):
+        engine_config = llm_config.get_engine_config()
+        _, vllm_config = _get_vllm_engine_config(llm_config)
+        self._model_config = vllm_config.model_config
+
+        frontend_args = FrontendArgs(**engine_config.frontend_kwargs)
+        self._renderer = OnlineRenderer(
+            self._model_config,
+            renderer_from_config(vllm_config),
+            request_logger=None,
+            chat_template=load_chat_template(frontend_args.chat_template),
+            chat_template_content_format=frontend_args.chat_template_content_format,
+            trust_request_chat_template=frontend_args.trust_request_chat_template,
+            default_chat_template_kwargs=frontend_args.default_chat_template_kwargs,
+        )
+        logger.info(
+            "In-process pre-routing tokenizer ready for %s",
+            self._model_config.model,
+        )
 
     async def tokenize(self, payload: Dict[str, Any]) -> Optional[List[int]]:
         """Tokenize a request ``payload`` into prompt token IDs.
@@ -49,76 +118,60 @@ async def tokenize(self, payload: Dict[str, Any]) -> Optional[List[int]]:
         Raises:
             TokenizeError: The ``/tokenize`` endpoint rejected the request.
         """
-        tok_req = self._build_tokenize_request(payload)
-        if tok_req is None:
+        request = build_tokenize_request(payload)
+        if request is None:
             return None
 
-        # /tokenize yields a single response; drain the stream fully so the
-        # handle response is cleaned up.
-        resp = None
-        async for chunk in self._handle.options(stream=True).tokenize.remote(
-            tok_req, None
-        ):
-            resp = chunk
-        if resp is None:
-            raise TokenizeError(
-                "/tokenize returned no response",
-                status_code=500,
-                type="internal_error",
-            )
-        if isinstance(resp, ErrorResponse):
-            raise TokenizeError(
-                resp.error.message,
-                status_code=resp.error.code,
-                type=resp.error.type,
-            )
-        return list(resp.tokens)
-
-    def _build_tokenize_request(
-        self, payload: Dict[str, Any]
-    ) -> Optional[Union[TokenizeChatRequest, TokenizeCompletionRequest]]:
-        """Build the Tokenize* request for ``payload``.
-
-        KV-aware routing sends each request to one replica, scored on a single
-        prompt's token sequence, so we return ``None`` (the caller falls back to
-        token-less routing) for bodies that don't have exactly one prompt:
-        - A non-string ``prompt``: an OpenAI *batch* completion where ``prompt``
-          is a list, e.g. ``{"prompt": ["q1", "q2"]}`` (or pre-tokenized id
-          lists). N prompts give N token sequences, so there's no single key to
-          route the one request on.
-
-        TODO (jeffreywang): Support multi-prompt tokenization.
-        """
         try:
-            if "messages" in payload:
-                # Forward every request field the engine renders the prompt from
-                # so the routing token IDs match the prefill tokens.
-                return TokenizeChatRequest.model_validate(
-                    {
-                        k: v
-                        for k, v in payload.items()
-                        if k in TokenizeChatRequest.model_fields
-                    }
+            if isinstance(request, TokenizeChatRequest):
+                engine_inputs = await self._render_chat(request)
+            else:
+                engine_inputs = await self._renderer.preprocess_completion(
+                    request,
+                    prompt_input=request.prompt,
+                    prompt_embeds=None,
+                    skip_mm_cache=True,
                 )
-            if "prompt" in payload:
-                if not isinstance(payload["prompt"], str):
-                    # TODO (jeffreywang): Multi-prompt (list) tokenization is unsupported;
-                    # fall back to token-less routing.
-                    return None
-                return TokenizeCompletionRequest.model_validate(
-                    {
-                        k: v
-                        for k, v in payload.items()
-                        if k in TokenizeCompletionRequest.model_fields
-                    }
-                )
-            # Should be unreachable: LLMRouter only routes bodies with messages
-            # or a prompt (see _parse_routing_payload).
-            logger.warning(
-                "Tokenizer got a payload with neither messages nor prompt; "
-                "falling back to token-less routing."
+        except TokenizeError:
+            raise
+        except (ValueError, jinja2.TemplateError) as e:
+            # /tokenize maps bad inputs and chat-template errors to 400; other
+            # exceptions are real bugs and should surface, not degrade routing.
+            raise TokenizeError(str(e), status_code=400, type="BadRequestError")
+
+        input_ids: List[int] = []
+        for engine_input in engine_inputs:
+            components = extract_prompt_components(self._model_config, engine_input)
+            if components.token_ids is not None:
+                input_ids.extend(components.token_ids)
+        return input_ids
+
+    async def _render_chat(self, request: TokenizeChatRequest):
+        # Refuse a request-supplied chat template unless the deployment opted in.
+        error = self._renderer.validate_chat_template(
+            request_chat_template=request.chat_template,
+            chat_template_kwargs=request.chat_template_kwargs,
+            trust_request_chat_template=self._renderer.trust_request_chat_template,
+        )
+        if error is not None:
+            raise TokenizeError(
+                error.error.message,
+                status_code=error.error.code,
+                type=error.error.type,
             )
-            return None
-        except Exception as e:
-            logger.debug("Unsupported tokenize request, falling back: %s", e)
-            return None
+
+        tool_dicts = (
+            None
+            if request.tools is None
+            else [tool.model_dump() for tool in request.tools]
+        )
+        _, engine_inputs = await self._renderer.preprocess_chat(
+            request,
+            request.messages,
+            default_template=self._renderer.chat_template,
+            default_template_content_format=self._renderer.chat_template_content_format,
+            default_template_kwargs=self._renderer.default_chat_template_kwargs,
+            tool_dicts=tool_dicts,
+            skip_mm_cache=True,
+        )
+        return engine_inputs

python/ray/llm/tests/serve/cpu/deployments/routers/kv/test_kv_aware_router.py

@@ -23,14 +23,14 @@
     LLMServingArgs,
     build_openai_app,
 )
+from ray.llm._internal.serve.routing_policies.kv_aware.constants import (
+    REQUEST_TOKEN_IDS_KWARG,
+)
 from ray.llm._internal.serve.routing_policies.kv_aware.kv_aware_actor import (
     KV_ROUTER_ACTOR_NAME,
     KVRouterActor,
     get_worker_id,
 )
-from ray.llm._internal.serve.routing_policies.kv_aware.tokenizer import (
-    REQUEST_TOKEN_IDS_KWARG,
-)
 from ray.serve._private.common import (
     REPLICA_ID_FULL_ID_STR_PREFIX,
     DeploymentID,