A high-throughput and memory-efficient inference and serving engine for LLMs
[Benchmark] Auto-detect and correct client/server tokenizer mismatch for random dataset (#44708)
akii96 committed
Jun 8, 2026 at 16:10 UTC
ac3409d162a92c4f57ccfd6ed37b7304b1e5747f
1 file changed
+81
-1
vllm/benchmarks/serve.py
+81
-1
@@ -31,7 +31,7 @@ import time
31
import uuid
32
import warnings
33
from collections.abc import AsyncGenerator, Iterable
34
-from dataclasses import dataclass
34
+from dataclasses import dataclass, replace
35
from datetime import datetime
36
from enum import Enum
37
from pathlib import Path
@@ -62,6 +62,80 @@ TERM_PLOTLIB_AVAILABLE = (importlib.util.find_spec("termplotlib") is not None) a
62
)
63
64
65
+async def _align_prompts_to_server_tokenizer(
66
+ base_url: str,
67
+ model_id: str,
68
+ input_requests: list[SampleRequest],
69
+ ssl_context: ssl.SSLContext | bool | None = None,
70
+) -> list[SampleRequest]:
71
+ """Re-align prompts if local/server tokenizers disagree."""
72
+ if not input_requests or not isinstance(input_requests[0].prompt, str):
73
+ return input_requests
74
+
75
+ tok_url = f"{base_url}/tokenize"
76
+ detok_url = f"{base_url}/detokenize"
77
+ connector = aiohttp.TCPConnector(ssl=ssl_context)
78
+
79
+ async with aiohttp.ClientSession(connector=connector) as session:
80
+ sem = asyncio.Semaphore(64)
81
+
82
+ async def _tokenize(prompt: str) -> list[int]:
83
+ async with (
84
+ sem,
85
+ session.post(
86
+ tok_url,
87
+ json={
88
+ "model": model_id,
89
+ "prompt": prompt,
90
+ "add_special_tokens": False,
91
+ },
92
+ ) as r,
93
+ ):
94
+ r.raise_for_status()
95
+ return (await r.json())["tokens"]
96
+
97
+ async def _detokenize(tokens: list[int]) -> str:
98
+ async with (
99
+ sem,
100
+ session.post(
101
+ detok_url, json={"model": model_id, "tokens": tokens}
102
+ ) as r,
103
+ ):
104
+ r.raise_for_status()
105
+ return (await r.json())["prompt"]
106
+
107
+ try:
108
+ first_tokens = await _tokenize(input_requests[0].prompt)
109
+ except Exception:
110
+ print("WARNING: /tokenize unavailable, skipping alignment.")
111
+ return input_requests
112
+
113
+ expected = input_requests[0].prompt_len
114
+ if len(first_tokens) == expected:
115
+ return input_requests
116
+
117
+ print(
118
+ f"WARNING: tokenizer mismatch "
119
+ f"(server={len(first_tokens)}, expected={expected}), "
120
+ f"re-aligning prompts."
121
+ )
122
+
123
+ async def _fix_one(req: SampleRequest) -> SampleRequest:
124
+ tokens = await _tokenize(req.prompt)
125
+ if len(tokens) <= req.prompt_len:
126
+ return req
127
+ corrected = await _detokenize(tokens[: req.prompt_len])
128
+ return replace(req, prompt=corrected, prompt_len=req.prompt_len)
129
+
130
+ results = await asyncio.gather(
131
+ *[_fix_one(r) for r in input_requests], return_exceptions=True
132
+ )
133
+ return [
134
+ res if not isinstance(res, BaseException) else orig
135
+ for orig, res in zip(input_requests, results)
136
+ ]
137
+
138
+
139
async def get_first_model_from_server(
140
base_url: str,
141
headers: dict | None = None,
@@ -1821,6 +1895,12 @@ async def main_async(args: argparse.Namespace) -> dict[str, Any]:
1895
1896
# Load the dataset.
1897
input_requests = get_samples(args, tokenizer)
1898
+
1899
+ if args.dataset_name in ("random", "prefix_repetition"):
1900
+ input_requests = await _align_prompts_to_server_tokenizer(
1901
+ base_url, model_id, input_requests, ssl_context
1902
+ )
1903
+
1904
goodput_config_dict = check_goodput_args(args)
1905
1906
backend = args.backend