From 954d971f8cd8626bb8e67c81e06dc2aa47223a6c Mon Sep 17 00:00:00 2001 From: Chenglun Hu Date: Sun, 16 Aug 2026 16:20:49 +0800 Subject: [PATCH] fix(scale-out): return all n choices from /inference/v1/generate serve_tokens set output_kind only for streaming, so non-streaming requests kept the SamplingParams default (CUMULATIVE). serve_tokens_full_generator keeps only the last streamed RequestOutput, and under CUMULATIVE each sequence is pruned from the stream once it finishes, so with n > 1 the sequences that finished before the final step never reach `choices` and the response silently returns fewer than n completions. The OpenAI completion path already uses FINAL_ONLY when not streaming; do the same here. Root cause and repro from @qgallouedec in #52398. Fixes #52398 Signed-off-by: Chenglun Hu --- .../token_in_token_out/test_serving_tokens.py | 21 +++++++++++++++++++ .../scale_out/token_in_token_out/serving.py | 10 +++++++-- 2 files changed, 29 insertions(+), 2 deletions(-) diff --git a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py index a1491c1c3e1b..cd0da519c1c7 100644 --- a/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py +++ b/tests/entrypoints/scale_out/token_in_token_out/test_serving_tokens.py @@ -462,3 +462,24 @@ async def test_generate_with_lora_adapter(client, tokenizer, messages): completions_res = completions_data["choices"][0]["message"]["content"] assert generate_res == completions_res + + +@pytest.mark.asyncio +@pytest.mark.parametrize("n", [2, 4]) +async def test_generate_endpoint_returns_all_n_choices(client, n): + # Regression test for #52398: a non-streaming /inference/v1/generate + # request with n > 1 must return exactly n choices. The handler previously + # left output_kind at the SamplingParams default (CUMULATIVE) for + # non-streaming requests, so sequences that finished before the final + # streamed step were dropped and fewer than n choices came back. + payload = { + "model": MODEL_NAME, + "token_ids": [9707, 1879, 11], + "sampling_params": {"max_tokens": 5, "temperature": 1.0, "n": n}, + "stream": False, + } + resp = await client.post(GEN_ENDPOINT, json=payload) + resp.raise_for_status() + data = resp.json() + assert len(data["choices"]) == n + assert sorted(c["index"] for c in data["choices"]) == list(range(n)) diff --git a/vllm/entrypoints/scale_out/token_in_token_out/serving.py b/vllm/entrypoints/scale_out/token_in_token_out/serving.py index 9e9ace877a7c..24978ffb3475 100644 --- a/vllm/entrypoints/scale_out/token_in_token_out/serving.py +++ b/vllm/entrypoints/scale_out/token_in_token_out/serving.py @@ -218,8 +218,14 @@ async def serve_tokens( if self.force_no_detokenize: sampling_params.detokenize = False - if request.stream: - sampling_params.output_kind = RequestOutputKind.DELTA + # Match the OpenAI serving path (CompletionRequest.to_sampling_params): + # non-streaming must use FINAL_ONLY. Under the SamplingParams default + # (CUMULATIVE) serve_tokens_full_generator keeps only the last streamed + # RequestOutput, and sequences that finished earlier are already pruned + # from it, so with n > 1 they silently never reach `choices`. + sampling_params.output_kind = ( + RequestOutputKind.DELTA if request.stream else RequestOutputKind.FINAL_ONLY + ) self._log_inputs( request_id,