A high-throughput and memory-efficient inference and serving engine for LLMs
[Bugfix] Fix --enable-prompt-tokens-details omitting zero cached tokens (#44383)
Signed-off-by: Sasindharan Sankar <sasindharansankar@email.com> Co-authored-by: Sasindharan Sankar <sasindharansankar@email.com> Co-authored-by: Chauncey <chaunceyjiang@gmail.com>
sasindharan committed
Jun 12, 2026 at 09:07 UTC
42ae5e7ac61910815bf368da22f67a721179ee45
5 files changed
+61
-9
tests/entrypoints/openai/completion/test_completion.py
+6
-3
@@ -58,9 +58,12 @@ async def test_single_completion(client: openai.AsyncOpenAI, model_name: str) ->
58
choice = completion.choices[0]
59
assert len(choice.text) >= 5
60
assert choice.finish_reason == "length"
61
- assert completion.usage == openai.types.CompletionUsage(
62
- completion_tokens=5, prompt_tokens=6, total_tokens=11
63
- )
61
+ assert completion.usage is not None
62
+ assert completion.usage.completion_tokens == 5
63
+ assert completion.usage.prompt_tokens == 6
64
+ assert completion.usage.total_tokens == 11
65
+ assert completion.usage.prompt_tokens_details is not None
66
+ assert completion.usage.prompt_tokens_details.cached_tokens == 0
67
68
# test using token IDs
69
completion = await client.completions.create(
tests/entrypoints/serve/disagg/test_generate_stream.py
+43
@@ -512,3 +512,46 @@ async def test_stream_prompt_tokens_details():
512
usage_chunk = parsed[-2]
513
assert usage_chunk["choices"] == []
514
assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 2
515
+
516
+
517
+@pytest.mark.asyncio
518
+async def test_stream_prompt_tokens_details_zero_cached():
519
+ """enable_prompt_tokens_details includes cached_tokens=0 in final usage.
520
+
521
+ Regression test for https://github.com/vllm-project/vllm/issues/44377:
522
+ zero cached tokens must not be treated as falsy and omitted.
523
+ """
524
+ engine = _mock_engine()
525
+
526
+ async def mock_generate(*args, **kwargs):
527
+ yield _make_request_output(
528
+ "req-1",
529
+ token_ids=[10],
530
+ finish_reason="stop",
531
+ finished=True,
532
+ num_cached_tokens=0,
533
+ )
534
+
535
+ engine.generate = MagicMock(side_effect=mock_generate)
536
+ serving = _build_serving_tokens(engine, enable_prompt_tokens_details=True)
537
+
538
+ request = GenerateRequest(
539
+ token_ids=[1, 2, 3],
540
+ sampling_params=SamplingParams(max_tokens=10),
541
+ model=MODEL_NAME,
542
+ stream=True,
543
+ stream_options=StreamOptions(include_usage=True),
544
+ )
545
+
546
+ response = await serving.serve_tokens(request)
547
+ chunks = []
548
+ async for chunk in response:
549
+ chunks.append(chunk)
550
+
551
+ parsed = _parse_sse_chunks(chunks)
552
+ # Usage-only chunk (before [DONE])
553
+ usage_chunk = parsed[-2]
554
+ assert usage_chunk["choices"] == []
555
+ # Zero cached tokens must be present, not omitted
556
+ assert usage_chunk["usage"]["prompt_tokens_details"] is not None
557
+ assert usage_chunk["usage"]["prompt_tokens_details"]["cached_tokens"] == 0
vllm/entrypoints/openai/chat_completion/serving.py
+5
-2
@@ -732,7 +732,7 @@ class OpenAIServingChat(OpenAIServing):
732
completion_tokens=completion_tokens,
733
total_tokens=num_prompt_tokens + completion_tokens,
734
)
735
- if self.enable_prompt_tokens_details and num_cached_tokens:
735
+ if self.enable_prompt_tokens_details and num_cached_tokens is not None:
736
final_usage.prompt_tokens_details = PromptTokenUsageInfo(
737
cached_tokens=num_cached_tokens
738
)
@@ -1023,7 +1023,10 @@ class OpenAIServingChat(OpenAIServing):
1023
completion_tokens=num_generated_tokens,
1024
total_tokens=num_prompt_tokens + num_generated_tokens,
1025
)
1026
- if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
1026
+ if (
1027
+ self.enable_prompt_tokens_details
1028
+ and final_res.num_cached_tokens is not None
1029
+ ):
1030
usage.prompt_tokens_details = PromptTokenUsageInfo(
1031
cached_tokens=final_res.num_cached_tokens
1032
)
vllm/entrypoints/openai/completion/serving.py
+2
-2
@@ -443,7 +443,7 @@ class OpenAIServingCompletion(OpenAIServing):
443
total_tokens=total_prompt_tokens + total_completion_tokens,
444
)
445
446
- if self.enable_prompt_tokens_details and num_cached_tokens:
446
+ if self.enable_prompt_tokens_details and num_cached_tokens is not None:
447
final_usage_info.prompt_tokens_details = PromptTokenUsageInfo(
448
cached_tokens=num_cached_tokens
449
)
@@ -583,7 +583,7 @@ class OpenAIServingCompletion(OpenAIServing):
583
if (
584
self.enable_prompt_tokens_details
585
and last_final_res
586
- and last_final_res.num_cached_tokens
586
+ and last_final_res.num_cached_tokens is not None
587
):
588
usage.prompt_tokens_details = PromptTokenUsageInfo(
589
cached_tokens=last_final_res.num_cached_tokens
vllm/entrypoints/serve/disagg/serving.py
+5
-2
@@ -307,7 +307,10 @@ class ServingTokens(OpenAIServing):
307
completion_tokens=num_generated_tokens,
308
total_tokens=num_prompt_tokens + num_generated_tokens,
309
)
310
- if self.enable_prompt_tokens_details and final_res.num_cached_tokens:
310
+ if (
311
+ self.enable_prompt_tokens_details
312
+ and final_res.num_cached_tokens is not None
313
+ ):
314
# This info is not available at the /coordinator level
315
usage.prompt_tokens_details = PromptTokenUsageInfo(
316
cached_tokens=final_res.num_cached_tokens
@@ -424,7 +427,7 @@ class ServingTokens(OpenAIServing):
427
total_tokens=num_prompt_tokens + total_completion_tokens,
428
)
429
427
- if self.enable_prompt_tokens_details and num_cached_tokens:
430
+ if self.enable_prompt_tokens_details and num_cached_tokens is not None:
431
final_usage_info.prompt_tokens_details = PromptTokenUsageInfo(
432
cached_tokens=num_cached_tokens
433
)