@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 25f2b55319
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 5, 14:45 UTC

[Frontend] feat: add streaming support for token generation endpoint (#37171)

Signed-off-by: Hyeonki Hong <hyeonki.hong@moreh.io>

Hyeonki Hong committed Apr 3, 2026 at 19:20 UTC 25f2b5531977293c780edabd211e826b9e644bff
4 files changed +665 -5
tests/entrypoints/serve/disagg/test_generate_stream.py new
+474
@@ -0,0 +1,474 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import json
5 +from dataclasses import dataclass, field
6 +from typing import Any
7 +from unittest.mock import AsyncMock, MagicMock
8 +
9 +import pytest
10 +
11 +from vllm.config.multimodal import MultiModalConfig
12 +from vllm.entrypoints.openai.engine.protocol import StreamOptions
13 +from vllm.entrypoints.openai.models.protocol import BaseModelPath
14 +from vllm.entrypoints.openai.models.serving import OpenAIServingModels
15 +from vllm.entrypoints.serve.disagg.protocol import GenerateRequest
16 +from vllm.entrypoints.serve.disagg.serving import ServingTokens
17 +from vllm.entrypoints.serve.render.serving import OpenAIServingRender
18 +from vllm.logprobs import Logprob
19 +from vllm.outputs import CompletionOutput, RequestOutput
20 +from vllm.renderers import renderer_from_config
21 +from vllm.sampling_params import SamplingParams
22 +from vllm.v1.engine.async_llm import AsyncLLM
23 +
24 +MODEL_NAME = "openai-community/gpt2"
25 +BASE_MODEL_PATHS = [
26 + BaseModelPath(name=MODEL_NAME, model_path=MODEL_NAME),
27 +]
28 +
29 +
30 +@dataclass
31 +class MockHFConfig:
32 + model_type: str = "any"
33 +
34 +
35 +@dataclass
36 +class MockModelConfig:
37 + task = "generate"
38 + runner_type = "generate"
39 + model = MODEL_NAME
40 + tokenizer = MODEL_NAME
41 + trust_remote_code = False
42 + tokenizer_mode = "auto"
43 + max_model_len = 100
44 + tokenizer_revision = None
45 + multimodal_config = MultiModalConfig()
46 + hf_config = MockHFConfig()
47 + hf_text_config = MockHFConfig()
48 + logits_processors: list[str] | None = None
49 + diff_sampling_param: dict | None = None
50 + allowed_local_media_path: str = ""
51 + allowed_media_domains: list[str] | None = None
52 + encoder_config = None
53 + generation_config: str = "auto"
54 + media_io_kwargs: dict[str, dict[str, Any]] = field(default_factory=dict)
55 + skip_tokenizer_init = False
56 + is_encoder_decoder: bool = False
57 + is_multimodal_model: bool = False
58 + renderer_num_workers: int = 1
59 +
60 + def get_diff_sampling_param(self):
61 + return self.diff_sampling_param or {}
62 +
63 +
64 +@dataclass
65 +class MockParallelConfig:
66 + _api_process_rank: int = 0
67 +
68 +
69 +@dataclass
70 +class MockVllmConfig:
71 + model_config: MockModelConfig
72 + parallel_config: MockParallelConfig
73 +
74 +
75 +def _build_renderer(model_config: MockModelConfig):
76 + return renderer_from_config(
77 + MockVllmConfig(model_config, parallel_config=MockParallelConfig()),
78 + )
79 +
80 +
81 +def _build_serving_tokens(engine: AsyncLLM, **kwargs) -> ServingTokens:
82 + models = OpenAIServingModels(
83 + engine_client=engine,
84 + base_model_paths=BASE_MODEL_PATHS,
85 + )
86 + serving_render = OpenAIServingRender(
87 + model_config=engine.model_config,
88 + renderer=engine.renderer,
89 + io_processor=engine.io_processor,
90 + model_registry=models.registry,
91 + request_logger=None,
92 + chat_template=None,
93 + chat_template_content_format="auto",
94 + )
95 + serving = ServingTokens(
96 + engine,
97 + models,
98 + openai_serving_render=serving_render,
99 + request_logger=None,
100 + **kwargs,
101 + )
102 +
103 + async def _fake_preprocess(*args, **kwargs):
104 + return [{"prompt_token_ids": [1, 2, 3]}]
105 +
106 + serving.openai_serving_render.preprocess_completion = AsyncMock(
107 + side_effect=_fake_preprocess
108 + )
109 + return serving
110 +
111 +
112 +def _make_request_output(
113 + request_id: str,
114 + token_ids: list[int],
115 + finish_reason: str | None = None,
116 + finished: bool = False,
117 + prompt_token_ids: list[int] | None = None,
118 + logprobs: list[dict[int, Any] | None] | None = None,
119 + num_cached_tokens: int | None = None,
120 + index: int = 0,
121 +) -> RequestOutput:
122 + return RequestOutput(
123 + request_id=request_id,
124 + prompt=None,
125 + prompt_token_ids=prompt_token_ids or [1, 2, 3],
126 + prompt_logprobs=None,
127 + outputs=[
128 + CompletionOutput(
129 + index=index,
130 + text="",
131 + token_ids=token_ids,
132 + cumulative_logprob=None,
133 + logprobs=logprobs,
134 + finish_reason=finish_reason,
135 + )
136 + ],
137 + finished=finished,
138 + metrics=None,
139 + lora_request=None,
140 + encoder_prompt=None,
141 + encoder_prompt_token_ids=None,
142 + num_cached_tokens=num_cached_tokens,
143 + )
144 +
145 +
146 +def _mock_engine() -> MagicMock:
147 + engine = MagicMock(spec=AsyncLLM)
148 + engine.errored = False
149 + engine.model_config = MockModelConfig()
150 + engine.input_processor = MagicMock()
151 + engine.io_processor = MagicMock()
152 + engine.renderer = _build_renderer(engine.model_config)
153 + return engine
154 +
155 +
156 +def _parse_sse_chunks(chunks: list[str]) -> list[Any]:
157 + """Parse SSE chunks into dicts (JSON) or raw strings ([DONE])."""
158 + parsed: list[Any] = []
159 + for chunk in chunks:
160 + assert chunk.startswith("data: ") and chunk.endswith("\n\n")
161 + payload = chunk[len("data: ") : -len("\n\n")]
162 + if payload == "[DONE]":
163 + parsed.append("[DONE]")
164 + else:
165 + parsed.append(json.loads(payload))
166 + return parsed
167 +
168 +
169 +@pytest.mark.asyncio
170 +async def test_stream_basic():
171 + """Streaming returns SSE chunks with correct token_ids and ends with [DONE]."""
172 + engine = _mock_engine()
173 +
174 + async def mock_generate(*args, **kwargs):
175 + yield _make_request_output("req-1", token_ids=[10])
176 + yield _make_request_output("req-1", token_ids=[20, 30])
177 + yield _make_request_output(
178 + "req-1", token_ids=[40], finish_reason="stop", finished=True
179 + )
180 +
181 + engine.generate = MagicMock(side_effect=mock_generate)
182 + serving = _build_serving_tokens(engine)
183 +
184 + request = GenerateRequest(
185 + token_ids=[1, 2, 3],
186 + sampling_params=SamplingParams(max_tokens=10),
187 + model=MODEL_NAME,
188 + stream=True,
189 + )
190 +
191 + response = await serving.serve_tokens(request)
192 + chunks = []
193 + async for chunk in response:
194 + chunks.append(chunk)
195 +
196 + parsed = _parse_sse_chunks(chunks)
197 +
198 + # 3 data chunks + [DONE]
199 + assert parsed[-1] == "[DONE]"
200 + data_chunks = [c for c in parsed if c != "[DONE]"]
201 + assert len(data_chunks) == 3
202 +
203 + assert data_chunks[0]["choices"][0]["token_ids"] == [10]
204 + assert data_chunks[1]["choices"][0]["token_ids"] == [20, 30]
205 + assert data_chunks[2]["choices"][0]["token_ids"] == [40]
206 + assert data_chunks[2]["choices"][0]["finish_reason"] == "stop"
207 +
208 +
209 +@pytest.mark.asyncio
210 +async def test_stream_error_mid_generation():
211 + """finish_reason='error' mid-stream yields error chunk then [DONE]."""
212 + engine = _mock_engine()
213 +
214 + async def mock_generate(*args, **kwargs):
215 + yield _make_request_output("req-1", token_ids=[10])
216 + yield _make_request_output(
217 + "req-1", token_ids=[20], finish_reason="error", finished=True
218 + )
219 +
220 + engine.generate = MagicMock(side_effect=mock_generate)
221 + serving = _build_serving_tokens(engine)
222 +
223 + request = GenerateRequest(
224 + token_ids=[1, 2, 3],
225 + sampling_params=SamplingParams(max_tokens=10),
226 + model=MODEL_NAME,
227 + stream=True,
228 + )
229 +
230 + response = await serving.serve_tokens(request)
231 + chunks = []
232 + async for chunk in response:
233 + chunks.append(chunk)
234 +
235 + assert len(chunks) >= 2
236 + assert any("Internal server error" in chunk for chunk in chunks), (
237 + f"Expected error message in chunks: {chunks}"
238 + )
239 + assert chunks[-1] == "data: [DONE]\n\n"
240 +
241 +
242 +@pytest.mark.asyncio
243 +async def test_stream_error_with_empty_delta():
244 + """finish_reason='error' with empty delta_token_ids still raises."""
245 + engine = _mock_engine()
246 +
247 + async def mock_generate(*args, **kwargs):
248 + yield _make_request_output("req-1", token_ids=[10])
249 + yield _make_request_output(
250 + "req-1", token_ids=[], finish_reason="error", finished=True
251 + )
252 +
253 + engine.generate = MagicMock(side_effect=mock_generate)
254 + serving = _build_serving_tokens(engine)
255 +
256 + request = GenerateRequest(
257 + token_ids=[1, 2, 3],
258 + sampling_params=SamplingParams(max_tokens=10),
259 + model=MODEL_NAME,
260 + stream=True,
261 + )
262 +
263 + response = await serving.serve_tokens(request)
264 + chunks = []
265 + async for chunk in response:
266 + chunks.append(chunk)
267 +
268 + assert any("Internal server error" in chunk for chunk in chunks), (
269 + f"Expected error message in chunks: {chunks}"
270 + )
271 + assert chunks[-1] == "data: [DONE]\n\n"
272 +
273 +
274 +@pytest.mark.asyncio
275 +async def test_stream_skips_empty_token_output():
276 + """Outputs with empty token_ids are skipped (no chunk emitted)."""
277 + engine = _mock_engine()
278 +
279 + async def mock_generate(*args, **kwargs):
280 + yield _make_request_output("req-1", token_ids=[10])
281 + yield _make_request_output("req-1", token_ids=[])
282 + yield _make_request_output(
283 + "req-1", token_ids=[20], finish_reason="stop", finished=True
284 + )
285 +
286 + engine.generate = MagicMock(side_effect=mock_generate)
287 + serving = _build_serving_tokens(engine)
288 +
289 + request = GenerateRequest(
290 + token_ids=[1, 2, 3],
291 + sampling_params=SamplingParams(max_tokens=10),
292 + model=MODEL_NAME,
293 + stream=True,
294 + )
295 +
296 + response = await serving.serve_tokens(request)
297 + chunks = []
298 + async for chunk in response:
299 + chunks.append(chunk)
300 +
301 + parsed = _parse_sse_chunks(chunks)
302 + assert parsed[-1] == "[DONE]"
303 + data_chunks = [c for c in parsed if c != "[DONE]"]
304 +
305 + # Only 2 data chunks — the empty one is skipped
306 + assert len(data_chunks) == 2
307 + assert data_chunks[0]["choices"][0]["token_ids"] == [10]
308 + assert data_chunks[1]["choices"][0]["token_ids"] == [20]
309 +
310 +
311 +@pytest.mark.asyncio
312 +async def test_stream_include_usage():
313 + """stream_options.include_usage emits a final usage-only chunk."""
314 + engine = _mock_engine()
315 +
316 + async def mock_generate(*args, **kwargs):
317 + yield _make_request_output("req-1", token_ids=[10])
318 + yield _make_request_output(
319 + "req-1", token_ids=[20], finish_reason="stop", finished=True
320 + )
321 +
322 + engine.generate = MagicMock(side_effect=mock_generate)
323 + serving = _build_serving_tokens(engine)
324 +
325 + request = GenerateRequest(
326 + token_ids=[1, 2, 3],
327 + sampling_params=SamplingParams(max_tokens=10),
328 + model=MODEL_NAME,
329 + stream=True,
330 + stream_options=StreamOptions(include_usage=True),
331 + )
332 +
333 + response = await serving.serve_tokens(request)
334 + chunks = []
335 + async for chunk in response:
336 + chunks.append(chunk)
337 +
338 + parsed = _parse_sse_chunks(chunks)
339 + assert parsed[-1] == "[DONE]"
340 +
341 + # The chunk before [DONE] should be the usage-only chunk
342 + usage_chunk = parsed[-2]
343 + assert usage_chunk["choices"] == []
344 + assert usage_chunk["usage"]["prompt_tokens"] == 3
345 + assert usage_chunk["usage"]["completion_tokens"] == 2
346 + assert usage_chunk["usage"]["total_tokens"] == 5
347 +
348 +
349 +@pytest.mark.asyncio
350 +async def test_stream_continuous_usage():
351 + """continuous_usage_stats adds usage to every data chunk."""
352 + engine = _mock_engine()
353 +
354 + async def mock_generate(*args, **kwargs):
355 + yield _make_request_output("req-1", token_ids=[10])
356 + yield _make_request_output(
357 + "req-1", token_ids=[20], finish_reason="stop", finished=True
358 + )
359 +
360 + engine.generate = MagicMock(side_effect=mock_generate)
361 + serving = _build_serving_tokens(engine)
362 +
363 + request = GenerateRequest(
364 + token_ids=[1, 2, 3],
365 + sampling_params=SamplingParams(max_tokens=10),
366 + model=MODEL_NAME,
367 + stream=True,
368 + stream_options=StreamOptions(
369 + include_usage=True,
370 + continuous_usage_stats=True,
371 + ),
372 + )
373 +
374 + response = await serving.serve_tokens(request)
375 + chunks = []
376 + async for chunk in response:
377 + chunks.append(chunk)
378 +
379 + parsed = _parse_sse_chunks(chunks)
380 + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
381 +
382 + # Every data chunk should have usage
383 + for i, dc in enumerate(data_chunks):
384 + assert dc["usage"] is not None, f"chunk {i} missing usage"
385 + assert dc["usage"]["prompt_tokens"] == 3
386 +
387 + # First chunk: 1 completion token
388 + assert data_chunks[0]["usage"]["completion_tokens"] == 1
389 + assert data_chunks[0]["usage"]["total_tokens"] == 4
390 +
391 + # Second chunk: 2 completion tokens (cumulative)
392 + assert data_chunks[1]["usage"]["completion_tokens"] == 2
393 + assert data_chunks[1]["usage"]["total_tokens"] == 5
394 +
395 +
396 +@pytest.mark.asyncio
397 +async def test_stream_with_logprobs():
398 + """Streaming with logprobs includes logprob data in each chunk."""
399 + engine = _mock_engine()
400 +
401 + async def mock_generate(*args, **kwargs):
402 + yield _make_request_output(
403 + "req-1",
404 + token_ids=[10],
405 + logprobs=[{10: Logprob(logprob=-0.5)}],
406 + )
407 + yield _make_request_output(
408 + "req-1",
409 + token_ids=[20],
410 + logprobs=[{20: Logprob(logprob=-1.0)}],
411 + finish_reason="stop",
412 + finished=True,
413 + )
414 +
415 + engine.generate = MagicMock(side_effect=mock_generate)
416 + serving = _build_serving_tokens(engine)
417 +
418 + request = GenerateRequest(
419 + token_ids=[1, 2, 3],
420 + sampling_params=SamplingParams(max_tokens=10, logprobs=1),
421 + model=MODEL_NAME,
422 + stream=True,
423 + )
424 +
425 + response = await serving.serve_tokens(request)
426 + chunks = []
427 + async for chunk in response:
428 + chunks.append(chunk)
429 +
430 + parsed = _parse_sse_chunks(chunks)
431 + data_chunks = [c for c in parsed if isinstance(c, dict) and c.get("choices")]
432 +
433 + for dc in data_chunks:
434 + lp = dc["choices"][0]["logprobs"]
435 + assert lp is not None
436 + assert len(lp["content"]) == 1
437 + assert lp["content"][0]["token"].startswith("token_id:")
438 +
439 +
440 +@pytest.mark.asyncio
441 +async def test_stream_prompt_tokens_details():
442 + """enable_prompt_tokens_details includes cached_tokens in final usage."""
443 + engine = _mock_engine()
444 +
445 + async def mock_generate(*args, **kwargs):
446 + yield _make_request_output(
447 + "req-1",
448 + token_ids=[10],
449 + finish_reason="stop",
450 + finished=True,
451 + num_cached_tokens=2,
452 + )
453 +
454 + engine.generate = MagicMock(side_effect=mock_generate)
455 + serving = _build_serving_tokens(engine, enable_prompt_tokens_details=True)
456 +
457 + request = GenerateRequest(
458 + token_ids=[1, 2, 3],
459 + sampling_params=SamplingParams(max_tokens=10),
460 + model=MODEL_NAME,
461 + stream=True,
462 + stream_options=StreamOptions(include_usage=True),
463 + )
464 +
465 + response = await serving.serve_tokens(request)
466 + chunks = []
467 + async for chunk in response:
468 + chunks.append(chunk)
469 +
470 + parsed = _parse_sse_chunks(chunks)
471 + # Usage-only chunk (before [DONE])
472 + usage_chunk = parsed[-2]
473 + assert usage_chunk["choices"] == []
474 + assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 2
tests/entrypoints/serve/disagg/test_serving_tokens.py
+49
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import json
5 import os
6
7 import httpx
@@ -113,6 +114,54 @@ async def test_generate_endpoint(client):
114 assert "choices" in data
115
116
117 +@pytest.mark.asyncio
118 +async def test_generate_stream(client):
119 + payload = {
120 + "model": MODEL_NAME,
121 + "token_ids": [1, 2, 3],
122 + "sampling_params": {"max_tokens": 5},
123 + "stream": True,
124 + }
125 + async with client.stream("POST", GEN_ENDPOINT, json=payload) as resp:
126 + resp.raise_for_status()
127 + chunks = []
128 + async for line in resp.aiter_lines():
129 + if not line.startswith("data: "):
130 + continue
131 + payload_str = line[len("data: ") :]
132 + if payload_str == "[DONE]":
133 + break
134 + chunks.append(json.loads(payload_str))
135 +
136 + assert len(chunks) > 0
137 + # Every chunk has choices with token_ids
138 + all_token_ids = []
139 + for chunk in chunks:
140 + assert "choices" in chunk
141 + assert len(chunk["choices"]) == 1
142 + choice = chunk["choices"][0]
143 + assert "token_ids" in choice
144 + assert len(choice["token_ids"]) > 0
145 + all_token_ids.extend(choice["token_ids"])
146 +
147 + # Last chunk should have a finish_reason
148 + assert chunks[-1]["choices"][0]["finish_reason"] is not None
149 +
150 + # Streaming should produce the same tokens as non-streaming
151 + non_stream_resp = await client.post(
152 + GEN_ENDPOINT,
153 + json={
154 + "model": MODEL_NAME,
155 + "token_ids": [1, 2, 3],
156 + "sampling_params": {"max_tokens": 5, "temperature": 0.0},
157 + "stream": False,
158 + },
159 + )
160 + non_stream_data = non_stream_resp.json()
161 + # Just verify we got the right number of tokens
162 + assert len(all_token_ids) == len(non_stream_data["choices"][0]["token_ids"])
163 +
164 +
165 @pytest.mark.asyncio
166 @pytest.mark.parametrize("logprobs_value", [0, 1, 5])
167 async def test_generate_logprobs(client, logprobs_value):
vllm/entrypoints/serve/disagg/protocol.py
+21 -1
@@ -6,7 +6,7 @@ from pydantic import BaseModel, Field, field_validator
6
7 from vllm.config import ModelConfig
8 from vllm.entrypoints.openai.chat_completion.protocol import ChatCompletionLogProbs
9 -from vllm.entrypoints.openai.engine.protocol import StreamOptions
9 +from vllm.entrypoints.openai.engine.protocol import StreamOptions, UsageInfo
10 from vllm.logprobs import Logprob
11 from vllm.renderers import TokenizeParams
12 from vllm.sampling_params import SamplingParams
@@ -122,6 +122,26 @@ class GenerateResponseChoice(BaseModel):
122 token_ids: list[int] | None = None
123
124
125 +class GenerateResponseStreamChoice(BaseModel):
126 + index: int
127 + logprobs: ChatCompletionLogProbs | None = None
128 + finish_reason: str | None = None
129 + token_ids: list[int] | None = None
130 +
131 +
132 +class GenerateStreamResponse(BaseModel):
133 + request_id: str = Field(
134 + default_factory=lambda: f"{random_uuid()}",
135 + description=(
136 + "The request_id related to this request. If the caller does "
137 + "not set it, a random_uuid will be generated. This id is used "
138 + "through out the inference process and return in response."
139 + ),
140 + )
141 + choices: list[GenerateResponseStreamChoice]
142 + usage: UsageInfo | None = Field(default=None)
143 +
144 +
145 class GenerateResponse(BaseModel):
146 request_id: str = Field(
147 default_factory=lambda: f"{random_uuid()}",
vllm/entrypoints/serve/disagg/serving.py
+121 -4
@@ -18,6 +18,7 @@ from vllm.entrypoints.openai.chat_completion.protocol import (
18 )
19 from vllm.entrypoints.openai.engine.protocol import (
20 ErrorResponse,
21 + GenerationError,
22 PromptTokenUsageInfo,
23 RequestResponseMetadata,
24 UsageInfo,
@@ -28,12 +29,15 @@ from vllm.entrypoints.serve.disagg.protocol import (
29 GenerateRequest,
30 GenerateResponse,
31 GenerateResponseChoice,
32 + GenerateResponseStreamChoice,
33 + GenerateStreamResponse,
34 )
35 from vllm.entrypoints.serve.render.serving import OpenAIServingRender
36 +from vllm.entrypoints.utils import should_include_usage
37 from vllm.logger import init_logger
38 from vllm.logprobs import Logprob
39 from vllm.outputs import RequestOutput
36 -from vllm.sampling_params import SamplingParams
40 +from vllm.sampling_params import RequestOutputKind, SamplingParams
41 from vllm.utils.collection_utils import as_list
42
43 logger = init_logger(__name__)
@@ -74,7 +78,7 @@ class ServingTokens(OpenAIServing):
78 self,
79 request: GenerateRequest,
80 raw_request: Request | None = None,
77 - ) -> GenerateResponse | ErrorResponse:
81 + ) -> GenerateResponse | ErrorResponse | AsyncGenerator[str, None]:
82 error_check_ret = await self._check_model(request)
83 if error_check_ret is not None:
84 logger.error("Error with model %s", error_check_ret)
@@ -110,6 +114,8 @@ class ServingTokens(OpenAIServing):
114 sampling_params = request.sampling_params
115 if self.force_no_detokenize:
116 sampling_params.detokenize = False
117 + if request.stream:
118 + sampling_params.output_kind = RequestOutputKind.DELTA
119
120 self._log_inputs(
121 request_id,
@@ -133,9 +139,17 @@ class ServingTokens(OpenAIServing):
139 priority=request.priority,
140 )
141
136 - # TODO(NickLucche): Implement streaming response
137 -
142 assert result_generator is not None
143 +
144 + if request.stream:
145 + return self.serve_tokens_stream_generator(
146 + request,
147 + result_generator,
148 + request_id,
149 + model_name,
150 + request_metadata,
151 + )
152 +
153 return await self.serve_tokens_full_generator(
154 request, result_generator, request_id, model_name, request_metadata
155 )
@@ -236,6 +250,109 @@ class ServingTokens(OpenAIServing):
250
251 return response
252
253 + async def serve_tokens_stream_generator(
254 + self,
255 + request: GenerateRequest,
256 + result_generator: AsyncGenerator[RequestOutput, None],
257 + request_id: str,
258 + model_name: str,
259 + request_metadata: RequestResponseMetadata,
260 + ) -> AsyncGenerator[str, None]:
261 + num_prompt_tokens = 0
262 + num_generated_tokens: list[int] = []
263 + first_iteration = True
264 + num_cached_tokens = None
265 + sampling_params: SamplingParams = request.sampling_params
266 +
267 + include_usage, include_continuous_usage = should_include_usage(
268 + request.stream_options, False
269 + )
270 +
271 + try:
272 + async for res in result_generator:
273 + if first_iteration:
274 + if res.prompt_token_ids is not None:
275 + num_prompt_tokens = len(res.prompt_token_ids)
276 + if res.encoder_prompt_token_ids is not None:
277 + num_prompt_tokens += len(res.encoder_prompt_token_ids)
278 + num_cached_tokens = res.num_cached_tokens
279 + num_generated_tokens = [0] * len(res.outputs)
280 + first_iteration = False
281 +
282 + for output in res.outputs:
283 + i = output.index
284 + delta_token_ids = output.token_ids
285 + num_generated_tokens[i] += len(delta_token_ids)
286 +
287 + finish_reason = output.finish_reason
288 + self._raise_if_error(finish_reason, request_id)
289 +
290 + if not delta_token_ids:
291 + continue
292 +
293 + if sampling_params.logprobs is not None:
294 + out_logprobs = output.logprobs
295 + assert out_logprobs is not None, "Did not output logprobs"
296 + logprobs = self._create_tokens_logprobs(
297 + token_ids=delta_token_ids,
298 + top_logprobs=out_logprobs,
299 + num_output_top_logprobs=sampling_params.logprobs,
300 + )
301 + else:
302 + logprobs = None
303 +
304 + chunk = GenerateStreamResponse(
305 + request_id=request_id,
306 + choices=[
307 + GenerateResponseStreamChoice(
308 + index=i,
309 + logprobs=logprobs,
310 + finish_reason=finish_reason,
311 + token_ids=as_list(delta_token_ids),
312 + )
313 + ],
314 + )
315 + if include_continuous_usage:
316 + chunk.usage = UsageInfo(
317 + prompt_tokens=num_prompt_tokens,
318 + completion_tokens=num_generated_tokens[i],
319 + total_tokens=(num_prompt_tokens + num_generated_tokens[i]),
320 + )
321 +
322 + yield f"data: {chunk.model_dump_json()}\n\n"
323 +
324 + total_completion_tokens = sum(num_generated_tokens)
325 + final_usage_info = UsageInfo(
326 + prompt_tokens=num_prompt_tokens,
327 + completion_tokens=total_completion_tokens,
328 + total_tokens=num_prompt_tokens + total_completion_tokens,
329 + )
330 +
331 + if self.enable_prompt_tokens_details and num_cached_tokens:
332 + final_usage_info.prompt_tokens_details = PromptTokenUsageInfo(
333 + cached_tokens=num_cached_tokens
334 + )
335 +
336 + if include_usage:
337 + final_chunk = GenerateStreamResponse(
338 + request_id=request_id,
339 + choices=[],
340 + usage=final_usage_info,
341 + )
342 + yield f"data: {final_chunk.model_dump_json(exclude_none=True)}\n\n"
343 +
344 + request_metadata.final_usage_info = final_usage_info
345 +
346 + except GenerationError as e:
347 + yield (
348 + f"data: {self._convert_generation_error_to_streaming_response(e)}\n\n"
349 + )
350 + except Exception as e:
351 + logger.exception("Error in token generation stream.")
352 + data = self.create_streaming_error_response(e)
353 + yield f"data: {data}\n\n"
354 + yield "data: [DONE]\n\n"
355 +
356 def _create_tokens_logprobs(
357 self,
358 token_ids: GenericSequence[int],