@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 6f1e7f7226
Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 13, 04:30 UTC

[DisaggEverything] Tokens in<>out `/generate` endpoint (#24261)

Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Nicolò Lucchesi committed Nov 14, 2025 at 17:58 UTC 6f1e7f7226447f606a0731376a2d0bd080aa2767
12 files changed +822 -9
examples/online_serving/token_generation_client.py new
+49
@@ -0,0 +1,49 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import httpx
4 +from transformers import AutoTokenizer
5 +
6 +GEN_ENDPOINT = "http://localhost:8000/inference/v1/generate"
7 +DUMMY_API_KEY = "empty"
8 +MODEL_NAME = "Qwen/Qwen3-0.6B"
9 +
10 +transport = httpx.HTTPTransport()
11 +headers = {"Authorization": f"Bearer {DUMMY_API_KEY}"}
12 +client = httpx.Client(
13 + transport=transport,
14 + base_url=GEN_ENDPOINT,
15 + timeout=600,
16 + headers=headers,
17 +)
18 +messages = [
19 + {"role": "system", "content": "You are a helpful assistant."},
20 + {"role": "user", "content": "How many countries are in the EU?"},
21 +]
22 +
23 +
24 +def main(client):
25 + tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
26 + token_ids = tokenizer.apply_chat_template(
27 + messages,
28 + add_generation_prompt=True,
29 + enable_thinking=False,
30 + )
31 + payload = {
32 + "model": MODEL_NAME,
33 + "token_ids": token_ids,
34 + "sampling_params": {"max_tokens": 24, "temperature": 0.2, "detokenize": False},
35 + "stream": False,
36 + }
37 + resp = client.post(GEN_ENDPOINT, json=payload)
38 + resp.raise_for_status()
39 + data = resp.json()
40 + print(data)
41 + print("-" * 50)
42 + print("Token generation results:")
43 + res = tokenizer.decode(data["choices"][0]["token_ids"])
44 + print(res)
45 + print("-" * 50)
46 +
47 +
48 +if __name__ == "__main__":
49 + main(client)
requirements/docs.txt
+4
@@ -10,3 +10,7 @@ mkdocs-minify-plugin
10 regex
11 ruff
12 pydantic
13 +
14 +# For generating argparse docs.
15 +# Adding requirements here should only be used as a last resort.
16 +msgspec # Need for multiple inheritance involving msgspec.Struct
\ No newline at end of file
tests/entrypoints/openai/test_serving_tokens.py new
+262
@@ -0,0 +1,262 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import httpx
5 +import pytest
6 +import pytest_asyncio
7 +from transformers import AutoTokenizer
8 +
9 +from vllm.config import ModelConfig
10 +from vllm.v1.engine.detokenizer import check_stop_strings
11 +
12 +from ...utils import RemoteOpenAIServer
13 +
14 +MODEL_NAME = "Qwen/Qwen3-0.6B"
15 +GEN_ENDPOINT = "/inference/v1/generate"
16 +
17 +
18 +def get_vocab_size(model_name):
19 + config = ModelConfig(
20 + model=model_name,
21 + seed=0,
22 + dtype="bfloat16",
23 + )
24 + return config.get_vocab_size()
25 +
26 +
27 +@pytest.fixture(scope="module")
28 +def tokenizer():
29 + return AutoTokenizer.from_pretrained(MODEL_NAME)
30 +
31 +
32 +@pytest.fixture(scope="module")
33 +def messages():
34 + return [
35 + {"role": "system", "content": "You are a helpful assistant."},
36 + {"role": "user", "content": "How many countries are in the EU?"},
37 + ]
38 +
39 +
40 +@pytest.fixture(scope="module")
41 +def server(request):
42 + args = [
43 + "--dtype",
44 + "bfloat16",
45 + "--max-model-len",
46 + "1024",
47 + "--enforce-eager",
48 + ]
49 +
50 + extra_args = getattr(request, "param", None)
51 + if extra_args is not None:
52 + args = args + (
53 + list(extra_args)
54 + if isinstance(extra_args, (list, tuple))
55 + else [str(extra_args)]
56 + )
57 +
58 + with RemoteOpenAIServer(MODEL_NAME, args) as remote_server:
59 + yield remote_server
60 +
61 +
62 +@pytest_asyncio.fixture
63 +async def client(server: RemoteOpenAIServer):
64 + transport = httpx.AsyncHTTPTransport(uds=server.uds) if server.uds else None
65 + headers = {"Authorization": f"Bearer {server.DUMMY_API_KEY}"}
66 + async with httpx.AsyncClient(
67 + transport=transport,
68 + base_url=server.url_root,
69 + timeout=600,
70 + headers=headers,
71 + ) as c:
72 + yield c
73 +
74 +
75 +@pytest.mark.asyncio
76 +async def test_generate_endpoint(client):
77 + payload = {
78 + "model": MODEL_NAME,
79 + "token_ids": [1, 2, 3],
80 + "sampling_params": {"max_tokens": 5},
81 + "stream": False,
82 + }
83 + resp = await client.post(GEN_ENDPOINT, json=payload)
84 + resp.raise_for_status()
85 + data = resp.json()
86 + assert "choices" in data
87 +
88 +
89 +@pytest.mark.asyncio
90 +async def test_same_response_as_chat_completions(client, tokenizer, messages):
91 + token_ids = tokenizer.apply_chat_template(
92 + messages,
93 + add_generation_prompt=True,
94 + enable_thinking=False, # default with Qwen3
95 + )
96 + for ignore_eos in [True, False]:
97 + payload = {
98 + "model": MODEL_NAME,
99 + "token_ids": token_ids,
100 + "sampling_params": {
101 + "max_tokens": 24,
102 + "temperature": 0.0,
103 + # NOTE coordinator will set this to skip detokenization
104 + "detokenize": False,
105 + "ignore_eos": ignore_eos,
106 + },
107 + "stream": False,
108 + }
109 + generate_resp = await client.post(GEN_ENDPOINT, json=payload)
110 + generate_data = generate_resp.json()
111 + generate_res = tokenizer.decode(
112 + generate_data["choices"][0]["token_ids"], skip_special_tokens=True
113 + )
114 +
115 + payload = {
116 + "model": MODEL_NAME,
117 + "messages": messages,
118 + "max_tokens": 24,
119 + "temperature": 0.0,
120 + "stream": False,
121 + "ignore_eos": ignore_eos,
122 + "chat_template_kwargs": dict(enable_thinking=False),
123 + }
124 + completions_resp = await client.post("/v1/chat/completions", json=payload)
125 + completions_data = completions_resp.json()
126 + completions_res = completions_data["choices"][0]["message"]["content"]
127 +
128 + assert generate_res == completions_res
129 +
130 +
131 +@pytest.mark.asyncio
132 +async def test_stop_string_workflow(client, tokenizer, messages):
133 + token_ids = tokenizer.apply_chat_template(
134 + messages,
135 + add_generation_prompt=True,
136 + enable_thinking=False, # default with Qwen3
137 + )
138 + payload = {
139 + "model": MODEL_NAME,
140 + "token_ids": token_ids,
141 + "sampling_params": {
142 + "max_tokens": 24,
143 + "temperature": 0.0,
144 + "detokenize": False,
145 + # stop strings are only supported when detokenize is True.
146 + "stop": ["27 member"],
147 + },
148 + # TODO stream test is much more interesting
149 + "stream": False,
150 + }
151 + with pytest.raises(httpx.HTTPStatusError):
152 + generate_resp = await client.post(GEN_ENDPOINT, json=payload)
153 + generate_resp.raise_for_status()
154 +
155 + payload["sampling_params"]["stop"] = None
156 + generate_resp = await client.post(
157 + GEN_ENDPOINT, json=payload, headers={"X-Request-Id": "42"}
158 + )
159 + generate_data = generate_resp.json()
160 + generate_res = tokenizer.decode(
161 + generate_data["choices"][0]["token_ids"], skip_special_tokens=True
162 + )
163 +
164 + # NOTE This is under the responsibility of the coordinator
165 + # stop_checker = StopChecker(
166 + # max_model_len=1024, get_tokenizer_for_seq=lambda _: tokenizer
167 + # )
168 + stop_str, truncate_to = check_stop_strings(
169 + generate_res, len(generate_res), ["27 member"], False
170 + )
171 + assert stop_str == "27 member"
172 + # abort request that hit stop string (requires tokens-only mode)
173 + # res = await client.post("/abort_requests", json={"request_ids": ["generate-tokens-42"]}) # noqa: E501
174 + # res.raise_for_status()
175 + generate_res = generate_res[:truncate_to]
176 +
177 + # Get stop_str response from chat completions
178 + payload = {
179 + "model": MODEL_NAME,
180 + "messages": messages,
181 + "max_tokens": 24,
182 + "temperature": 0.0,
183 + "stream": False,
184 + "stop": ["27 member"],
185 + "chat_template_kwargs": dict(enable_thinking=False),
186 + }
187 + completions_resp = await client.post("/v1/chat/completions", json=payload)
188 + completions_data = completions_resp.json()
189 + completions_res = completions_data["choices"][0]["message"]["content"]
190 + assert generate_res == completions_res
191 +
192 +
193 +@pytest.mark.asyncio
194 +@pytest.mark.parametrize(
195 + "server",
196 + [
197 + [
198 + "--enable-lora",
199 + "--lora-modules",
200 + "Alice=charent/self_cognition_Alice",
201 + "Bob=charent/self_cognition_Bob",
202 + "--max-lora-rank",
203 + "64",
204 + "--max-cpu-loras",
205 + "2",
206 + ]
207 + ],
208 + indirect=True,
209 +)
210 +async def test_generate_with_lora_adapter(client, tokenizer, messages):
211 + # Verify adapters are listed
212 + models_resp = await client.get("/v1/models")
213 + models_resp.raise_for_status()
214 + models = {m["id"] for m in models_resp.json().get("data", [])}
215 + assert {"Alice", "Bob"}.issubset(models)
216 +
217 + # Generate using a LoRA adapter by specifying its name as the model
218 + payload = {
219 + "model": "Alice",
220 + "token_ids": [1, 2, 3],
221 + "sampling_params": {"max_tokens": 5},
222 + "stream": False,
223 + }
224 + resp = await client.post(GEN_ENDPOINT, json=payload)
225 + resp.raise_for_status()
226 + data = resp.json()
227 + assert "choices" in data
228 +
229 + token_ids = tokenizer.apply_chat_template(
230 + messages,
231 + add_generation_prompt=True,
232 + enable_thinking=False, # default with Qwen3
233 + )
234 + payload = {
235 + "model": "Alice",
236 + "token_ids": token_ids,
237 + "sampling_params": {
238 + "max_tokens": 24,
239 + "temperature": 0.0,
240 + "detokenize": False,
241 + },
242 + "stream": False,
243 + }
244 + generate_resp = await client.post(GEN_ENDPOINT, json=payload)
245 + generate_data = generate_resp.json()
246 + generate_res = tokenizer.decode(
247 + generate_data["choices"][0]["token_ids"], skip_special_tokens=True
248 + )
249 +
250 + payload = {
251 + "model": "Alice",
252 + "messages": messages,
253 + "max_tokens": 24,
254 + "temperature": 0.0,
255 + "stream": False,
256 + "chat_template_kwargs": dict(enable_thinking=False),
257 + }
258 + completions_resp = await client.post("/v1/chat/completions", json=payload)
259 + completions_data = completions_resp.json()
260 + completions_res = completions_data["choices"][0]["message"]["content"]
261 +
262 + assert generate_res == completions_res
vllm/engine/arg_utils.py
+5
@@ -566,6 +566,7 @@ class EngineArgs:
566 kv_offloading_backend: KVOffloadingBackend | None = (
567 CacheConfig.kv_offloading_backend
568 )
569 + tokens_only: bool = False
570
571 def __post_init__(self):
572 # support `EngineArgs(compilation_config={...})`
@@ -1495,6 +1496,10 @@ class EngineArgs:
1496 else ParallelConfig.data_parallel_rpc_port
1497 )
1498
1499 + if self.tokens_only and not model_config.skip_tokenizer_init:
1500 + model_config.skip_tokenizer_init = True
1501 + logger.info("Skipping tokenizer initialization for tokens-only mode.")
1502 +
1503 # Forward the deprecated CLI args to the EPLB config.
1504 if self.num_redundant_experts is not None:
1505 self.eplb_config.num_redundant_experts = self.num_redundant_experts
vllm/entrypoints/openai/api_server.py
+81
@@ -65,6 +65,8 @@ from vllm.entrypoints.openai.protocol import (
65 EmbeddingResponse,
66 ErrorInfo,
67 ErrorResponse,
68 + GenerateRequest,
69 + GenerateResponse,
70 IOProcessorResponse,
71 PoolingBytesResponse,
72 PoolingRequest,
@@ -96,6 +98,7 @@ from vllm.entrypoints.openai.serving_pooling import OpenAIServingPooling
98 from vllm.entrypoints.openai.serving_responses import OpenAIServingResponses
99 from vllm.entrypoints.openai.serving_score import ServingScores
100 from vllm.entrypoints.openai.serving_tokenization import OpenAIServingTokenization
101 +from vllm.entrypoints.openai.serving_tokens import ServingTokens
102 from vllm.entrypoints.openai.serving_transcription import (
103 OpenAIServingTranscription,
104 OpenAIServingTranslation,
@@ -357,6 +360,10 @@ def engine_client(request: Request) -> EngineClient:
360 return request.app.state.engine_client
361
362
363 +def generate_tokens(request: Request) -> ServingTokens | None:
364 + return request.app.state.serving_tokens
365 +
366 +
367 @router.get("/health", response_class=Response)
368 async def health(raw_request: Request) -> Response:
369 """Health check."""
@@ -1228,6 +1235,41 @@ INVOCATION_VALIDATORS = [
1235 ]
1236
1237
1238 +@router.post(
1239 + "/inference/v1/generate",
1240 + dependencies=[Depends(validate_json_request)],
1241 + responses={
1242 + HTTPStatus.OK.value: {"content": {"text/event-stream": {}}},
1243 + HTTPStatus.BAD_REQUEST.value: {"model": ErrorResponse},
1244 + HTTPStatus.NOT_FOUND.value: {"model": ErrorResponse},
1245 + HTTPStatus.INTERNAL_SERVER_ERROR.value: {"model": ErrorResponse},
1246 + },
1247 +)
1248 +@with_cancellation
1249 +@load_aware_call
1250 +async def generate(request: GenerateRequest, raw_request: Request):
1251 + handler = generate_tokens(raw_request)
1252 + if handler is None:
1253 + return base(raw_request).create_error_response(
1254 + message="The model does not support generate tokens API"
1255 + )
1256 + try:
1257 + generator = await handler.serve_tokens(request, raw_request)
1258 + except Exception as e:
1259 + raise HTTPException(
1260 + status_code=HTTPStatus.INTERNAL_SERVER_ERROR.value, detail=str(e)
1261 + ) from e
1262 + if isinstance(generator, ErrorResponse):
1263 + return JSONResponse(
1264 + content=generator.model_dump(), status_code=generator.error.code
1265 + )
1266 +
1267 + elif isinstance(generator, GenerateResponse):
1268 + return JSONResponse(content=generator.model_dump())
1269 +
1270 + return StreamingResponse(content=generator, media_type="text/event-stream")
1271 +
1272 +
1273 if envs.VLLM_TORCH_PROFILER_DIR:
1274 logger.warning_once(
1275 "Torch Profiler is enabled in the API server. This should ONLY be "
@@ -1629,6 +1671,31 @@ def build_app(args: Namespace) -> FastAPI:
1671 )
1672
1673 app = sagemaker_standards.bootstrap(app)
1674 + # Optional endpoints
1675 + if args.tokens_only:
1676 +
1677 + @app.post("/abort_requests")
1678 + async def abort_requests(raw_request: Request):
1679 + """
1680 + Abort one or more requests. To be used in a
1681 + Disaggregated Everything setup.
1682 + """
1683 + try:
1684 + body = await raw_request.json()
1685 + except json.JSONDecodeError as e:
1686 + raise HTTPException(
1687 + status_code=HTTPStatus.BAD_REQUEST.value,
1688 + detail=f"JSON decode error: {e}",
1689 + ) from e
1690 + request_ids = body.get("request_ids")
1691 + if request_ids is None:
1692 + raise HTTPException(
1693 + status_code=HTTPStatus.BAD_REQUEST.value,
1694 + detail="Missing 'request_ids' in request body",
1695 + )
1696 + # Abort requests in background
1697 + asyncio.create_task(engine_client(raw_request).abort(request_ids))
1698 + return Response(status_code=200)
1699
1700 return app
1701
@@ -1851,6 +1918,20 @@ async def init_app_state(
1918 if "generate" in supported_tasks
1919 else None
1920 )
1921 + state.serving_tokens = (
1922 + ServingTokens(
1923 + engine_client,
1924 + state.openai_serving_models,
1925 + request_logger=request_logger,
1926 + return_tokens_as_token_ids=args.return_tokens_as_token_ids,
1927 + log_error_stack=args.log_error_stack,
1928 + enable_prompt_tokens_details=args.enable_prompt_tokens_details,
1929 + enable_log_outputs=args.enable_log_outputs,
1930 + force_no_detokenize=args.tokens_only,
1931 + )
1932 + if "generate" in supported_tasks
1933 + else None
1934 + )
1935
1936 state.enable_server_load_tracking = args.enable_server_load_tracking
1937 state.server_load_metrics = 0
vllm/entrypoints/openai/cli_args.py
+5
@@ -189,6 +189,11 @@ class FrontendArgs:
189 Helps mitigate header abuse. Default: 256."""
190 log_error_stack: bool = envs.VLLM_SERVER_DEV_MODE
191 """If set to True, log the stack trace of error responses"""
192 + tokens_only: bool = False
193 + """
194 + If set to True, only enable the Tokens In<>Out endpoint.
195 + This is intended for use in a Disaggregated Everything setup.
196 + """
197
198 @staticmethod
199 def add_cli_args(parser: FlexibleArgumentParser) -> FlexibleArgumentParser:
vllm/entrypoints/openai/protocol.py
+77
@@ -3220,3 +3220,80 @@ class TranslationResponseVerbose(OpenAIBaseModel):
3220
3221 words: list[TranslationWord] | None = None
3222 """Extracted words and their corresponding timestamps."""
3223 +
3224 +
3225 +####### Tokens IN <> Tokens OUT #######
3226 +class GenerateRequest(BaseModel):
3227 + request_id: str = Field(
3228 + default_factory=lambda: f"{random_uuid()}",
3229 + description=(
3230 + "The request_id related to this request. If the caller does "
3231 + "not set it, a random_uuid will be generated. This id is used "
3232 + "through out the inference process and return in response."
3233 + ),
3234 + )
3235 + token_ids: list[int]
3236 + """The token ids to generate text from."""
3237 +
3238 + # features: MultiModalFeatureSpec
3239 + # TODO (NickLucche): implement once Renderer work is completed
3240 + features: str | None = None
3241 + """The processed MM inputs for the model."""
3242 +
3243 + sampling_params: SamplingParams
3244 + """The sampling parameters for the model."""
3245 +
3246 + model: str | None = None
3247 +
3248 + stream: bool | None = False
3249 + stream_options: StreamOptions | None = None
3250 + cache_salt: str | None = Field(
3251 + default=None,
3252 + description=(
3253 + "If specified, the prefix cache will be salted with the provided "
3254 + "string to prevent an attacker to guess prompts in multi-user "
3255 + "environments. The salt should be random, protected from "
3256 + "access by 3rd parties, and long enough to be "
3257 + "unpredictable (e.g., 43 characters base64-encoded, corresponding "
3258 + "to 256 bit)."
3259 + ),
3260 + )
3261 + priority: int = Field(
3262 + default=0,
3263 + description=(
3264 + "The priority of the request (lower means earlier handling; "
3265 + "default: 0). Any priority other than 0 will raise an error "
3266 + "if the served model does not use priority scheduling."
3267 + ),
3268 + )
3269 + kv_transfer_params: dict[str, Any] | None = Field(
3270 + default=None,
3271 + description="KVTransfer parameters used for disaggregated serving.",
3272 + )
3273 +
3274 +
3275 +class GenerateResponseChoice(BaseModel):
3276 + index: int
3277 + logprobs: ChatCompletionLogProbs | None = None
3278 + # per OpenAI spec this is the default
3279 + finish_reason: str | None = "stop"
3280 + token_ids: list[int] | None = None
3281 +
3282 +
3283 +class GenerateResponse(BaseModel):
3284 + request_id: str = Field(
3285 + default_factory=lambda: f"{random_uuid()}",
3286 + description=(
3287 + "The request_id related to this request. If the caller does "
3288 + "not set it, a random_uuid will be generated. This id is used "
3289 + "through out the inference process and return in response."
3290 + ),
3291 + )
3292 + choices: list[GenerateResponseChoice]
3293 +
3294 + prompt_logprobs: list[dict[int, Logprob] | None] | None = None
3295 +
3296 + kv_transfer_params: dict[str, Any] | None = Field(
3297 + default=None,
3298 + description="KVTransfer parameters used for disaggregated serving.",
3299 + )
vllm/entrypoints/openai/serving_engine.py
+4
@@ -58,6 +58,8 @@ from vllm.entrypoints.openai.protocol import (
58 ErrorResponse,
59 FunctionCall,
60 FunctionDefinition,
61 + GenerateRequest,
62 + GenerateResponse,
63 IOProcessorRequest,
64 PoolingResponse,
65 RerankRequest,
@@ -134,6 +136,7 @@ AnyRequest: TypeAlias = (
136 | SpeechToTextRequest
137 | ResponsesRequest
138 | IOProcessorRequest
139 + | GenerateRequest
140 )
141
142 AnyResponse: TypeAlias = (
@@ -145,6 +148,7 @@ AnyResponse: TypeAlias = (
148 | PoolingResponse
149 | ClassificationResponse
150 | ScoreResponse
151 + | GenerateResponse
152 )
153
154
vllm/entrypoints/openai/serving_tokens.py new
+269
@@ -0,0 +1,269 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import asyncio
4 +import time
5 +from collections.abc import AsyncGenerator
6 +from collections.abc import Sequence as GenericSequence
7 +
8 +from fastapi import Request
9 +
10 +# yapf: disable
11 +from vllm.engine.protocol import EngineClient
12 +from vllm.entrypoints.logger import RequestLogger
13 +from vllm.entrypoints.openai.protocol import (
14 + ChatCompletionLogProb,
15 + ChatCompletionLogProbs,
16 + ChatCompletionLogProbsContent,
17 + ErrorResponse,
18 + GenerateRequest,
19 + GenerateResponse,
20 + GenerateResponseChoice,
21 + PromptTokenUsageInfo,
22 + RequestResponseMetadata,
23 + UsageInfo,
24 +)
25 +from vllm.entrypoints.openai.serving_engine import OpenAIServing, clamp_prompt_logprobs
26 +from vllm.entrypoints.openai.serving_models import OpenAIServingModels
27 +from vllm.inputs.data import TokensPrompt as EngineTokensPrompt
28 +from vllm.logger import init_logger
29 +from vllm.logprobs import Logprob
30 +from vllm.outputs import RequestOutput
31 +from vllm.sampling_params import SamplingParams
32 +from vllm.utils.collection_utils import as_list
33 +
34 +logger = init_logger(__name__)
35 +
36 +
37 +class ServingTokens(OpenAIServing):
38 + """Provides Tokens IN <> Tokens OUT functionality to vLLM API."""
39 +
40 + def __init__(
41 + self,
42 + engine_client: EngineClient,
43 + models: OpenAIServingModels,
44 + *,
45 + request_logger: RequestLogger | None,
46 + force_no_detokenize: bool = False,
47 + return_tokens_as_token_ids: bool = False,
48 + log_error_stack: bool = False,
49 + enable_prompt_tokens_details: bool = False,
50 + enable_log_outputs: bool = False,
51 + ):
52 + super().__init__(engine_client=engine_client,
53 + models=models,
54 + request_logger=request_logger,
55 + return_tokens_as_token_ids=return_tokens_as_token_ids,
56 + log_error_stack=log_error_stack)
57 + self.enable_prompt_tokens_details = enable_prompt_tokens_details
58 + self.enable_log_outputs = enable_log_outputs
59 + self.force_no_detokenize = force_no_detokenize
60 + if force_no_detokenize:
61 + logger.info("Tokens-only mode is enabled, skipping detokenization "
62 + "step for incoming requests.")
63 +
64 + async def serve_tokens(
65 + self,
66 + request: GenerateRequest,
67 + raw_request: Request | None = None
68 + ) -> GenerateResponse | ErrorResponse:
69 + error_check_ret = await self._check_model(request)
70 + if error_check_ret is not None:
71 + logger.error("Error with model %s", error_check_ret)
72 + return error_check_ret
73 +
74 + # If the engine is dead, raise the engine's DEAD_ERROR.
75 + # This is required for the streaming case, where we return a
76 + # success status before we actually start generating text :).
77 + if self.engine_client.errored:
78 + raise self.engine_client.dead_error
79 +
80 + lora_request = None
81 + lora_request = self._maybe_get_adapters(request,
82 + supports_default_mm_loras=True)
83 +
84 + model_name = self.models.model_name(lora_request)
85 +
86 + request_id = "generate-tokens-" \
87 + f"{self._base_request_id(raw_request, request.request_id)}"
88 +
89 + request_metadata = RequestResponseMetadata(request_id=request_id)
90 + if raw_request:
91 + raw_request.state.request_metadata = request_metadata
92 +
93 + # TODO(NickLucche): Change to EngineCoreRequest once Renderer work is
94 + # completed
95 + engine_prompt = EngineTokensPrompt(prompt_token_ids=request.token_ids)
96 + if request.features is not None:
97 + engine_prompt["multi_modal_data"] = None
98 +
99 + if hasattr(request, "cache_salt") and request.cache_salt is not None:
100 + engine_prompt["cache_salt"] = request.cache_salt
101 +
102 + # Schedule the request and get the result generator.
103 + result_generator: AsyncGenerator[RequestOutput, None] | None = None
104 + try:
105 + sampling_params = request.sampling_params
106 + if self.force_no_detokenize:
107 + sampling_params.detokenize = False
108 +
109 + self._log_inputs(request_id,
110 + request.token_ids,
111 + params=sampling_params,
112 + lora_request=lora_request)
113 +
114 + trace_headers = (None if raw_request is None else await
115 + self._get_trace_headers(raw_request.headers))
116 +
117 + result_generator = self.engine_client.generate(
118 + engine_prompt,
119 + sampling_params,
120 + request_id,
121 + lora_request=lora_request,
122 + trace_headers=trace_headers,
123 + priority=request.priority,
124 + )
125 +
126 + except ValueError as e:
127 + return self.create_error_response(str(e))
128 +
129 + # TODO(NickLucche): Implement streaming response
130 +
131 + try:
132 + assert result_generator is not None
133 + return await self.serve_tokens_full_generator(
134 + request, result_generator, request_id, model_name,
135 + request_metadata)
136 + except ValueError as e:
137 + return self.create_error_response(str(e))
138 +
139 + async def serve_tokens_full_generator(
140 + self,
141 + request: GenerateRequest,
142 + result_generator: AsyncGenerator[RequestOutput, None],
143 + request_id: str,
144 + model_name: str,
145 + request_metadata: RequestResponseMetadata,
146 + ) -> ErrorResponse | GenerateResponse:
147 +
148 + created_time = int(time.time())
149 + final_res: RequestOutput | None = None
150 + sampling_params: SamplingParams = request.sampling_params
151 +
152 + try:
153 + async for res in result_generator:
154 + final_res = res
155 + except asyncio.CancelledError:
156 + return self.create_error_response("Client disconnected")
157 + except ValueError as e:
158 + return self.create_error_response(str(e))
159 +
160 + assert final_res is not None
161 +
162 + choices: list[GenerateResponseChoice] = []
163 + num_generated_tokens = 0
164 + for output in final_res.outputs:
165 + token_ids = output.token_ids
166 + out_logprobs = output.logprobs
167 +
168 + # This is top_logprobs in completions API
169 + if sampling_params.logprobs:
170 + assert out_logprobs is not None, "Did not output logprobs"
171 + logprobs = self._create_tokens_logprobs(
172 + token_ids=token_ids,
173 + top_logprobs=out_logprobs,
174 + num_output_top_logprobs=sampling_params.logprobs,
175 + )
176 + else:
177 + logprobs = None
178 +
179 + choice_data = GenerateResponseChoice(
180 + index=output.index,
181 + logprobs=logprobs,
182 + finish_reason=output.finish_reason
183 + if output.finish_reason else "stop",
184 + token_ids=as_list(output.token_ids))
185 +
186 + choices.append(choice_data)
187 + num_generated_tokens += len(output.token_ids)
188 +
189 + assert final_res.prompt_token_ids is not None
190 + num_prompt_tokens = len(final_res.prompt_token_ids)
191 + if final_res.encoder_prompt_token_ids is not None:
192 + num_prompt_tokens += len(final_res.encoder_prompt_token_ids)
193 +
194 + usage = UsageInfo(prompt_tokens=num_prompt_tokens,
195 + completion_tokens=num_generated_tokens,
196 + total_tokens=num_prompt_tokens +
197 + num_generated_tokens)
198 + if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
199 + # This info is not available at the /coordinator level
200 + usage.prompt_tokens_details = PromptTokenUsageInfo(
201 + cached_tokens=final_res.num_cached_tokens)
202 +
203 + request_metadata.final_usage_info = usage
204 +
205 + response = GenerateResponse(
206 + id=request_id,
207 + created=created_time,
208 + model=model_name,
209 + choices=choices,
210 + usage=usage,
211 + prompt_logprobs=clamp_prompt_logprobs(final_res.prompt_logprobs),
212 + kv_transfer_params=final_res.kv_transfer_params,
213 + )
214 +
215 + # Log complete response if output logging is enabled
216 + if self.enable_log_outputs and self.request_logger:
217 + for choice in choices:
218 + # Get the corresponding output token IDs
219 + output_token_ids = None
220 + if choice.index < len(final_res.outputs):
221 + output_token_ids = final_res.outputs[
222 + choice.index].token_ids
223 +
224 + if output_token_ids:
225 + # Log token_ids only.
226 + self.request_logger.log_outputs(
227 + request_id=request_id,
228 + outputs="",
229 + output_token_ids=output_token_ids,
230 + finish_reason=choice.finish_reason,
231 + is_streaming=False,
232 + delta=False,
233 + )
234 +
235 + return response
236 +
237 + def _create_tokens_logprobs(
238 + self,
239 + token_ids: GenericSequence[int],
240 + top_logprobs: GenericSequence[dict[int, Logprob] | None],
241 + num_output_top_logprobs: int | None = None,
242 + ) -> ChatCompletionLogProbs:
243 + """Create OpenAI-style logprobs."""
244 + logprobs_content: list[ChatCompletionLogProbsContent] = []
245 +
246 + for i, token_id in enumerate(token_ids):
247 + token = f"token_id:{token_id}"
248 + step_top_logprobs = top_logprobs[i]
249 + if step_top_logprobs is None or step_top_logprobs.get(
250 + token_id) is None:
251 + logprobs_content.append(
252 + ChatCompletionLogProbsContent(token=token, ))
253 + else:
254 + step_token = step_top_logprobs[token_id]
255 +
256 + logprobs_content.append(
257 + ChatCompletionLogProbsContent(
258 + token=token,
259 + logprob=max(step_token.logprob, -9999.0),
260 + top_logprobs=[
261 + ChatCompletionLogProb(
262 + token=token,
263 + logprob=max(p[1].logprob, -9999.0),
264 + ) for i, p in enumerate(step_top_logprobs.items())
265 + if num_output_top_logprobs
266 + and i < num_output_top_logprobs
267 + ]))
268 +
269 + return ChatCompletionLogProbs(content=logprobs_content)
vllm/sampling_params.py
+2
@@ -15,6 +15,7 @@ from pydantic.dataclasses import dataclass
15 from vllm.logger import init_logger
16 from vllm.logits_process import LogitsProcessor
17 from vllm.transformers_utils.tokenizer import AnyTokenizer
18 +from vllm.v1.serial_utils import PydanticMsgspecMixin
19
20 logger = init_logger(__name__)
21
@@ -122,6 +123,7 @@ class RequestOutputKind(Enum):
123
124
125 class SamplingParams(
126 + PydanticMsgspecMixin,
127 msgspec.Struct,
128 omit_defaults=True, # type: ignore[call-arg]
129 # required for @cached_property.
vllm/v1/engine/__init__.py
+1 -7
@@ -15,6 +15,7 @@ from vllm.pooling_params import PoolingParams
15 from vllm.sampling_params import SamplingParams
16 from vllm.v1.metrics.stats import SchedulerStats
17 from vllm.v1.outputs import LogprobsLists, LogprobsTensors
18 +from vllm.v1.serial_utils import UtilityResult
19
20 # These are possible values of RequestOutput.finish_reason,
21 # so form part of the external API.
@@ -131,13 +132,6 @@ class EngineCoreOutput(
132 return self.finish_reason is not None
133
134
134 -class UtilityResult:
135 - """Wrapper for special handling when serializing/deserializing."""
136 -
137 - def __init__(self, r: Any = None):
138 - self.result = r
139 -
140 -
135 class UtilityOutput(
136 msgspec.Struct,
137 array_like=True, # type: ignore[call-arg]
vllm/v1/serial_utils.py
+63 -2
@@ -8,7 +8,7 @@ from collections.abc import Callable, Sequence
8 from functools import partial
9 from inspect import isclass
10 from types import FunctionType
11 -from typing import Any, TypeAlias
11 +from typing import Any, TypeAlias, get_type_hints
12
13 import cloudpickle
14 import msgspec
@@ -16,6 +16,8 @@ import numpy as np
16 import torch
17 import zmq
18 from msgspec import msgpack
19 +from pydantic import GetCoreSchemaHandler
20 +from pydantic_core import core_schema
21
22 from vllm import envs
23 from vllm.logger import init_logger
@@ -32,7 +34,6 @@ from vllm.multimodal.inputs import (
34 NestedTensors,
35 )
36 from vllm.utils.platform_utils import is_pin_memory_available
35 -from vllm.v1.engine import UtilityResult
37 from vllm.v1.utils import tensor_data
38
39 logger = init_logger(__name__)
@@ -104,6 +105,13 @@ def _decode_type_info_recursive(
105 return convert_fn(type_info, data)
106
107
108 +class UtilityResult:
109 + """Wrapper for special handling when serializing/deserializing."""
110 +
111 + def __init__(self, r: Any = None):
112 + self.result = r
113 +
114 +
115 class MsgpackEncoder:
116 """Encoder with custom torch tensor and numpy array serialization.
117
@@ -469,3 +477,56 @@ def run_method(
477 else:
478 func = partial(method, obj) # type: ignore
479 return func(*args, **kwargs)
480 +
481 +
482 +class PydanticMsgspecMixin:
483 + @classmethod
484 + def __get_pydantic_core_schema__(
485 + cls, source_type: Any, handler: GetCoreSchemaHandler
486 + ) -> core_schema.CoreSchema:
487 + """
488 + Make msgspec.Struct compatible with Pydantic, respecting defaults.
489 + Handle JSON=>msgspec.Struct. Used when exposing msgspec.Struct to the
490 + API as input or in `/docs`. Note this is cached by Pydantic and not
491 + called on every validation.
492 + """
493 + msgspec_fields = {f.name: f for f in msgspec.structs.fields(source_type)}
494 + type_hints = get_type_hints(source_type)
495 +
496 + # Build the Pydantic typed_dict_field for each msgspec field
497 + fields = {}
498 + for name, hint in type_hints.items():
499 + msgspec_field = msgspec_fields[name]
500 +
501 + # typed_dict_field using the handler to get the schema
502 + field_schema = handler(hint)
503 +
504 + # Add default value to the schema.
505 + if msgspec_field.default_factory is not msgspec.NODEFAULT:
506 + wrapped_schema = core_schema.with_default_schema(
507 + schema=field_schema,
508 + default_factory=msgspec_field.default_factory,
509 + )
510 + fields[name] = core_schema.typed_dict_field(wrapped_schema)
511 + elif msgspec_field.default is not msgspec.NODEFAULT:
512 + wrapped_schema = core_schema.with_default_schema(
513 + schema=field_schema,
514 + default=msgspec_field.default,
515 + )
516 + fields[name] = core_schema.typed_dict_field(wrapped_schema)
517 + else:
518 + # No default, so Pydantic will treat it as required
519 + fields[name] = core_schema.typed_dict_field(field_schema)
520 + return core_schema.no_info_after_validator_function(
521 + cls._validate_msgspec,
522 + core_schema.typed_dict_schema(fields),
523 + )
524 +
525 + @classmethod
526 + def _validate_msgspec(cls, value: Any) -> Any:
527 + """Validate and convert input to msgspec.Struct instance."""
528 + if isinstance(value, cls):
529 + return value
530 + if isinstance(value, dict):
531 + return cls(**value)
532 + return msgspec.convert(value, type=cls)