Fix full logprobs in token-in/token-out responses (#58488)
Signed-off-by: aoshen02 <aoshen@inferact.ai>
aoshen02 committed
Sep 24, 2026 at 12:46 UTC
7f1a5398e9610d96c473931a26c0e12bbe0d0423
2 files changed
+14
-1
tests/entrypoints/scale_out/token_in_token_out/test_tokens_logprobs.py
+10
@@ -26,3 +26,13 @@ def test_logprobs_zero_emits_sampled_token():
26
num_output_top_logprobs=0,
27
)
28
assert len(result.content[0].top_logprobs) == 1
29
+
30
+
31
+def test_logprobs_minus_one_emits_all_tokens():
32
+ result = ServingTokens._create_tokens_logprobs(
33
+ None,
34
+ token_ids=[7],
35
+ top_logprobs=[{7: Logprob(-0.9), 8: Logprob(-1.1)}],
36
+ num_output_top_logprobs=-1,
37
+ )
38
+ assert len(result.content[0].top_logprobs) == 2
vllm/entrypoints/scale_out/token_in_token_out/serving.py
+4
-1
@@ -590,7 +590,10 @@ class ServingTokens(GenerateBaseServing):
590
step_top_logprobs.items()
591
)
592
if num_output_top_logprobs is not None
593
- and i < max(num_output_top_logprobs, 1)
593
+ and (
594
+ num_output_top_logprobs == -1
595
+ or i < max(num_output_top_logprobs, 1)
596
+ )
597
],
598
)
599
)