[Feature] Add return_token_ids to /v1/chat/completions and /v1/completions

[ianliuy]

Motivation

Closes #18378

For RL training workflows, token consistency between inference and training requires the exact token IDs — retokenizing decoded text may not reproduce the same IDs. The /generate endpoint already returns output_ids; this PR brings the same capability to the OpenAI-compatible chat and completions endpoints.

Note: PR #18398 covers the same feature and has been open for 2+ months with no reviews. This is a fresh, minimal implementation that addresses correctness concerns and follows codebase conventions more closely.

Changes

protocol.py

• ChatCompletionRequest + CompletionRequest: add return_token_ids: bool = False (opt-in, off by default)
• SglExt: add prompt_token_ids: Optional[List[int]] and completion_token_ids: Optional[List[List[int]]]
  • _serialize already strips None fields → zero wire-format impact when disabled

serving_chat.py

• Streaming: accumulate output_ids per choice (plain overwrite — output_ids is cumulative like routed_experts), emit one choices=[] sglext chunk after the stream loop
• Non-streaming: read prompt_token_ids from adapted_request.input_ids (guarded for multimodal), completion_token_ids from ret items

serving_completions.py

• Same pattern as serving_chat.py
• Mirrors existing symmetry: all return_* fields appear in both CompletionRequest and ChatCompletionRequest

test/registered/openai_server/basic/test_protocol.py

• 7 new unit tests in TestReturnTokenIds

Usage

response = client.chat.completions.create(
    model="...",
    messages=[{"role": "user", "content": "Hello!"}],
    extra_body={"return_token_ids": True},
)
sglext = response.model_extra.get("sglext", {})
print(sglext["prompt_token_ids"])       # [1, 2, 3, ...]
print(sglext["completion_token_ids"])   # [[4, 5, 6, ...]]  (one list per choice)

Design notes

• Zero overhead when disabled (all logic gated on return_token_ids)
• Non-OAI-standard fields placed in sglext per existing convention
• Multimodal: prompt_token_ids returns None (nested input_ids not supported)
• Backend unchanged — tokenizer_manager.py already returns output_ids in every chunk

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

[gemini-code-assist]

Code Review

This pull request introduces the return_token_ids feature to the OpenAI-compatible Chat and Completion endpoints, allowing users to receive prompt and completion token IDs within the sglext response field. The implementation supports both streaming and non-streaming modes and includes new unit tests for the protocol changes. Review feedback highlights that prompt_token_ids should be returned even if no output tokens are generated. Furthermore, the completions endpoint requires updates to include prompt_token_ids in streaming chunks and to handle string-based prompts by tokenizing them to ensure the requested IDs are correctly populated in the response.

[gemini-code-assist]

The check and output_ids_accum prevents returning prompt_token_ids if no output tokens were generated (e.g., when max_tokens=0 or the model returns an empty response). Since prompt_token_ids are available regardless of the generation outcome, this check should be removed to ensure the requested data is always returned.

Suggested change

	if request.return_token_ids and output_ids_accum:
	if request.return_token_ids:

[gemini-code-assist]

This streaming chunk is missing the prompt_token_ids field, which creates a discrepancy with the chat endpoint. Additionally, the and output_ids_accum check should be removed to ensure prompt_token_ids are returned even if no output tokens were generated. For completions where the prompt is a string, we should also attempt to retrieve the token IDs from the adapted request's text.

            if request.return_token_ids:
                n_choices = request.n or 1
                prompt_ids = getattr(adapted_request, "input_ids", None)
                if prompt_ids is None:
                    text = getattr(adapted_request, "text", None)
                    if isinstance(text, str):
                        prompt_ids = self.tokenizer_manager.tokenizer.encode(text)

                # Only include prompt_ids for non-multimodal requests (flat list of ints)
                if (
                    isinstance(prompt_ids, list)
                    and prompt_ids
                    and isinstance(prompt_ids[0], list)
                ):
                    prompt_ids = None

                completion_token_ids = [
                    output_ids_accum.get(i, []) for i in range(n_choices)
                ]
                token_ids_chunk = CompletionStreamResponse(
                    id=content["meta_info"]["id"],
                    created=created,
                    object="text_completion",
                    choices=[],
                    model=request.model,
                    sglext=SglExt(
                        completion_token_ids=completion_token_ids,
                        prompt_token_ids=prompt_ids,
                    ),
                )
                yield f"data: {token_ids_chunk.model_dump_json()}\n\n"

[gemini-code-assist]

For completions requests where the prompt is provided as a string (the most common case), adapted_request.input_ids will be None. This results in prompt_token_ids being omitted from the response even when requested. The prompt should be tokenized if input_ids is missing to ensure the feature works as expected for string prompts.

            input_ids = getattr(adapted_request, "input_ids", None)
            if input_ids is None:
                text = getattr(adapted_request, "text", None)
                if isinstance(text, str):
                    input_ids = self.tokenizer_manager.tokenizer.encode(text)

            if (
                input_ids is not None
                and len(input_ids) > 0
                and not isinstance(input_ids[0], list)
            ):
                prompt_token_ids = list(input_ids)

[MLKoz2]

Up

[MLKoz2]

What do you think about change completion_token_ids to token_ids as in vLLM?
To turn it on we turn on return_token_ids so token_ids is more intuitive and coherent with effort-less serving engines changes.

[anerli]

+1 would love to get this functionality.

However, this should support token IDs for streamed SSE chunks as well.
Each SSE chunk in vLLM returns token_ids, like:

{"id":"chatcmpl-abc","object":"chat.completion.chunk","created":1750000000,"model":"my-model","choices":[{"index":0,"delta": {"content":"Hello"},"logprobs":null,"finish_reason":null,"token_ids":[9707]}]}

If the streamed token IDs were added in a way that follows the vLLM convention, it may be logical to also follow the convention for the overall completion tokens as @MLKoz2 suggested.

For reference, vLLM PR: vllm-project/vllm#22587