[A5bhinav]

Refs #682

/v1/completions serves n from 1 to 128: each prompt is fanned out into n ordinary engine requests, and their outputs come back as n choices, streamed or not. The engine is unchanged, and chat still refuses n != 1. Builds on #842, which this replaces the refusal of; until that merges, the diff here also shows its commit, and the change in this PR is the last one. Opened as a draft: it has not been run on a GPU yet.

Summary

After #842, completions refuse n != 1. This PR makes them serve it, following OpenAI's conventions: choices are ordered prompt-major (index = prompt_index * n + j); usage.prompt_tokens counts each prompt once and completion_tokens is summed over all choices; when streaming, every chunk carries its choice index, each choice gets exactly one finish_reason chunk, then one usage chunk (with include_usage) and a single [DONE]; n < 1 or n > 128 (OpenAI's own ceiling) is a 400 with param "n".

• server/generation.py: ChoiceFanOut, a protocol-neutral helper that runs n generations of one prompt and merges their outputs as (choice_index, item) pairs, so chat can reuse it later.
  • Choice 0 is submitted first and the other n - 1 follow its first sampled token. By then its prompt is in the radix prefix cache (cache_req(finished=False) in the scheduler's _process_last_data runs before that token's reply is sent), so they should get a prefix hit instead of prefilling the same prompt again. The engine's first reply only reports admission and comes before the cache insert, so it does not release them. This is from reading the scheduler; I have not measured it on hardware.
  • All n uids are allocated up front in start(), so admission is decided once: a server that starts stopping cannot fail a request halfway through, and an admission error is a 503 before a stream opens.
  • If any choice fails, or the consumer stops early (error, disconnect, cancellation), every choice not yet finished is aborted, including uids never submitted, which clears their stats entries. On a choice-0 error the siblings are never sent.
• server/openai_api.py: handle_completion builds a ChoiceFanOut per prompt (_completion_fan_out) for both paths; stream_completion_chunks tags chunks with the choice index and emits one finish chunk per choice. A streaming response that never starts aborts choice 0 from the response's background task. For n = 1 the only visible difference is that a streamed completion's id is now "cmpl-" plus a random hex string, like the non-streamed one, instead of "cmpl-" plus the uid.
• server/api_server.py: stream_with_cancellation also accepts a callable returning the uids to abort, since a fan-out's uids are not all live when the stream starts.

Open choices I would like a maintainer's view on: the cap of 128 (an alternative is tying it to --max-running-requests); cached_tokens is reported from choice 0 only, so it does not show the reuse by choices 1..n-1; prompts in a list still run one after another, as before, and only the n choices of one prompt run concurrently.

Known limitations: Starlette skips a streaming response's background task on a client disconnect under ASGI spec >= 2.4, so a stream that never starts leaves choice 0 running there (uvicorn, which FreeToken runs on, reports 2.3; the same gap exists for n = 1 on main). A non-streaming client that disconnects does not cancel its generations, also unchanged from main. Each abort of a never-submitted uid leaves a scheduler tombstone, at most n - 1 per request that ends early; the scheduler already caps them at 65,536.

Why now: clients that pass n to completions (best-of selection, self-consistency, eval harnesses) otherwise have to send n separate requests, each prefilling the same prompt. Fanning out on the server keeps the OpenAI request shape and prefills the prompt once.

Tests

In tests/server/test_openai_api.py, with a small fake state (MultiUserState) that hands out one reply list per uid and logs sends, replies and aborts. test_completion_refuses_n_other_than_one from #842 becomes test_completion_refuses_n_out_of_range, for n in {0, 129}, streamed and not. New:

• test_completion_n_orders_choices_and_counts_prompt_once: choices are ordered prompt-major across a list of prompts, and each prompt is counted once in usage;
• test_completion_n_submits_siblings_after_first_sampled_token: siblings are sent only after choice 0's first sampled token, not after its admission reply;
• test_completion_n_error_aborts_sibling_choices: an error in one choice returns a 400 and aborts the others;
• test_completion_n_choice_zero_error_submits_no_siblings: an error on choice 0 sends no siblings;
• test_stream_completion_n_tags_chunks_and_finishes_each_choice_once: streamed chunks carry the right index, each choice finishes exactly once, and usage is correct.

Each of these fails if the behaviour it covers is removed.

Test environment

• GPU: none. I don't have access to an NVIDIA GPU yet, so this has not been run against a live ft serve.
• CPU: Apple M4 Pro, 24 GB RAM, macOS 26.6.2
• torch 2.14.1 (CPU), Python 3.12.13
• This branch at 662c266, A/B against main e328cf6.

Unit tests

PYTHONPATH=python pytest tests/server/test_openai_api.py

This branch: 20 passed, 17 failed, 1 error. Main: 11 passed, 17 failed, 1 error. The 9 extra passes are the tests above. The 17 failures and the error are all "ModuleNotFoundError: No module named 'flashlib'", which needs CUDA, and are the same on main.

End-to-end A/B against main

Not run yet: it needs a GPU, which I don't have access to yet. Once I do, I will run these and add the results here before taking this out of draft:

• the repro attached to /v1/completions accepts n and silently returns one choice; /v1/chat/completions refuses n != 1 with 400 #682 (python repro_completions_n_ignored.py --base-url http://127.0.0.1:1941): on this branch every row should PASS, with choices [0, 1] / [0, 1, 2];
• n=4 at temperature 1.0 should give 4 different texts;
• one n=4 request on this branch against 4 concurrent n=1 requests on main, same long prompt: TTFT, total time, and #cached-token in the prefill log for choices 1..n-1.