From 4b6e6778bf454a66d5ecd5947c8ead6da0896ef4 Mon Sep 17 00:00:00 2001 From: Asthenia Date: Sat, 29 Aug 2026 10:09:27 +0000 Subject: [PATCH 1/2] fix: Use per-model tokenizer in KvawareRouter for multi-model setups KvawareRouter kept a single self.tokenizer that was initialized from endpoints[0].model_names[0] on the first request and reused for every later request. When the router serves multiple base models, requests for model B were tokenized with model A's tokenizer, causing incorrect token IDs to be sent to the LMCache KV-aware lookup. Replace the single tokenizer with a per-model cache (self.tokenizers: Dict[str, AutoTokenizer]) so that each model gets its own tokenizer. The remote /tokenize fallback also uses the correct endpoint for the requested model. Closes #1052 Signed-off-by: Asthenia --- src/vllm_router/routers/routing_logic.py | 24 +++++++++++++++++------- 1 file changed, 17 insertions(+), 7 deletions(-) diff --git a/src/vllm_router/routers/routing_logic.py b/src/vllm_router/routers/routing_logic.py index b25e9184a..f339c88e9 100644 --- a/src/vllm_router/routers/routing_logic.py +++ b/src/vllm_router/routers/routing_logic.py @@ -326,7 +326,7 @@ def __init__( self.instance_id_to_ip = {} self.session_key = session_key self.hash_ring = HashRing() - self.tokenizer = None + self.tokenizers: Dict[str, AutoTokenizer] = {} self.threshold = kv_aware_threshold def start_kv_manager(self): @@ -388,18 +388,28 @@ async def route_request( token_ids = None # Local-first tokenization, fall back to remote "/tokenize" API on failure # TODO (Yuhan): Handle chat completions + model_name = request_json.get("model", "") + # Find the endpoint serving this model + model_endpoint = next( + (ep for ep in endpoints if model_name in ep.model_names), + endpoints[0] if endpoints else None, + ) try: - if self.tokenizer is None: - self.tokenizer = AutoTokenizer.from_pretrained( - endpoints[0].model_names[0] + if model_name not in self.tokenizers: + self.tokenizers[model_name] = AutoTokenizer.from_pretrained( + model_name ) - token_ids = self.tokenizer.encode(request_json.get("prompt", "")) + token_ids = self.tokenizers[model_name].encode( + request_json.get("prompt", "") + ) except Exception: # Remote /tokenize fallback (let errors bubble up to keep behavior simple) - remote_url = endpoints[0].url + "/tokenize" + if model_endpoint is None: + raise + remote_url = model_endpoint.url + "/tokenize" headers = {"Content-Type": "application/json"} data = { - "model": endpoints[0].model_names[0], + "model": model_name, "prompt": request_json.get("prompt", ""), } body = requests.post( From 31c010916c8f5cb9199e95104043bec46551ad34 Mon Sep 17 00:00:00 2001 From: Asthenia Date: Sat, 29 Aug 2026 10:12:05 +0000 Subject: [PATCH 2/2] =?UTF-8?q?fix:=20Address=20review=20feedback=20?= =?UTF-8?q?=E2=80=94=20fix=20tokenize=5Fprompt=20and=20model=20name=20hand?= =?UTF-8?q?ling?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Fix tokenize_prompt() used by LoadAwareRouter to use per-model tokenizer cache instead of the removed self.tokenizer - Add validation for missing model name in request body - Refactor KvawareRouter.route_request() to reuse tokenize_prompt() instead of inline tokenization, avoiding code duplication Signed-off-by: Asthenia --- src/vllm_router/routers/routing_logic.py | 53 +++++++++--------------- 1 file changed, 19 insertions(+), 34 deletions(-) diff --git a/src/vllm_router/routers/routing_logic.py b/src/vllm_router/routers/routing_logic.py index f339c88e9..c6a5cca48 100644 --- a/src/vllm_router/routers/routing_logic.py +++ b/src/vllm_router/routers/routing_logic.py @@ -388,34 +388,7 @@ async def route_request( token_ids = None # Local-first tokenization, fall back to remote "/tokenize" API on failure # TODO (Yuhan): Handle chat completions - model_name = request_json.get("model", "") - # Find the endpoint serving this model - model_endpoint = next( - (ep for ep in endpoints if model_name in ep.model_names), - endpoints[0] if endpoints else None, - ) - try: - if model_name not in self.tokenizers: - self.tokenizers[model_name] = AutoTokenizer.from_pretrained( - model_name - ) - token_ids = self.tokenizers[model_name].encode( - request_json.get("prompt", "") - ) - except Exception: - # Remote /tokenize fallback (let errors bubble up to keep behavior simple) - if model_endpoint is None: - raise - remote_url = model_endpoint.url + "/tokenize" - headers = {"Content-Type": "application/json"} - data = { - "model": model_name, - "prompt": request_json.get("prompt", ""), - } - body = requests.post( - remote_url, headers=headers, json=data, timeout=10 - ).json() - token_ids = body["tokens"] + token_ids = await self.tokenize_prompt(endpoints, request_json) event_id = "Lookup" + str(uuid.uuid4()) msg = LookupMsg(tokens=token_ids, event_id=event_id) @@ -662,17 +635,29 @@ async def tokenize_prompt( The remote fallback is a blocking HTTP call, so it runs in an executor rather than on the event loop. """ + model_name = request_json.get("model", "") + if not model_name: + raise ValueError("Missing 'model' in request body") + # Find the endpoint serving this model + model_endpoint = next( + (ep for ep in endpoints if model_name in ep.model_names), + endpoints[0] if endpoints else None, + ) try: - if self.tokenizer is None: - self.tokenizer = AutoTokenizer.from_pretrained( - endpoints[0].model_names[0] + if model_name not in self.tokenizers: + self.tokenizers[model_name] = AutoTokenizer.from_pretrained( + model_name ) - return self.tokenizer.encode(request_json.get("prompt", "")) + return self.tokenizers[model_name].encode( + request_json.get("prompt", "") + ) except Exception: - remote_url = endpoints[0].url + "/tokenize" + if model_endpoint is None: + raise + remote_url = model_endpoint.url + "/tokenize" headers = {"Content-Type": "application/json"} data = { - "model": endpoints[0].model_names[0], + "model": model_name, "prompt": request_json.get("prompt", ""), } loop = asyncio.get_running_loop()