[OAI Server Refactor] [ChatCompletions & Completions] Support Return Hidden State

[key4ng]

Motivation

Address issue: #7260
Support return hidden states in serving_chat.py and serving_completions.py

• Complete task below
• Add UTs

Modifications

• Add enable_hidden_state field for request and hiddle_state field for response in protocol.py
• Add hidden_state logic in serving_chat.py and serving_completions.py
• Added unit test

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Summary of Changes

Hello @key4ng, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces the capability for users to request and receive the hidden states (embeddings) from the last token of a model's output for both Chat Completions and standard Completions API calls. The changes involve updating the API protocol, modifying the serving logic to handle and return these states, and adding a dedicated utility for processing them, all backed by a new suite of unit tests to ensure correctness and robustness.

Highlights

• API Protocol Extension: Extended the OpenAI-compatible API protocol (protocol.py) to include a new return_hidden_states boolean field in CompletionRequest and ChatCompletionRequest models. Correspondingly, a hidden_states field (Optional[object]) has been added to CompletionResponseChoice, CompletionResponseStreamChoice, ChatCompletionResponseChoice, and DeltaMessage to carry the model's hidden states in the response.
• Hidden State Processing Logic: Implemented logic within serving_chat.py and serving_completions.py to process and return hidden states. This includes passing the return_hidden_states flag to the internal request, extracting hidden states from the model's output, and including the last token's hidden states in both non-streaming and streaming API responses. For streaming, hidden states are sent as a separate chunk.
• Utility Function for Hidden States: Added a new utility function, process_hidden_states_from_ret, in utils.py to standardize the extraction and formatting of hidden states from the internal response structure, ensuring only the last token's hidden states are returned when requested.
• Comprehensive Unit Tests: Introduced extensive unit tests in test_serving_chat.py and test_serving_completions.py to validate the new hidden states functionality across various scenarios, including single/multiple requests, streaming/non-streaming responses, echo behavior, multiple choices, and edge cases like empty or single-token hidden states.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist is currently in preview and may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This PR adds support for returning hidden states in chat and completion endpoints, with updates to protocol definitions and serving logic. The changes are well-tested. Suggestions focus on enhancing type specificity in the protocol and simplifying conditional logic.

[gemini-code-assist]

Simplify the condition for extracting last_token_hidden_states by removing the redundant choice_hidden_states check.

                            last_token_hidden_states = (
                                choice_hidden_states[-1] if len(choice_hidden_states) > 1 else []
                            )

[CatherineSue]

This suggestion is valid. There is a if choice_hidden_states: at L625.

[CatherineSue]

I feel like the original implementation in adapter.py is cleaner, it doesn't put everything inside a
if choice_hidden_states.

What is the reason behind this change?

Original link:
https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/openai_api/adapter.py#L1849-L1870

[key4ng]

thanks, yeah line 625 can be removed

[key4ng]

my concern is if choice_hidden_states is none or empty, do we send the hidden states chunk? In the implementation of adapter.py, it will still send
nvm, exclude_if_none can handle that

[CatherineSue]

This suggestion is valid. There is a if choice_hidden_states: at L625.

[CatherineSue]

I feel like the original implementation in adapter.py is cleaner, it doesn't put everything inside a
if choice_hidden_states.

What is the reason behind this change?

Original link:
https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/openai_api/adapter.py#L1849-L1870

[key4ng]

/gemini summary

[gemini-code-assist]

Summary of Changes

This pull request enhances the OpenAI server's capabilities by allowing clients to retrieve the hidden states from the language model's output. This feature is integrated into both chat and completion API endpoints, providing deeper insights into the model's internal representations. The changes involve extending the request and response protocols, implementing the necessary data processing within the serving logic, and adding robust unit tests to ensure correctness.

Highlights

• API Protocol Extension: Introduced a new return_hidden_states boolean field in CompletionRequest and ChatCompletionRequest to allow users to explicitly request the model's hidden states. Corresponding hidden_states fields (as Optional[object]) were added to the various response choice and delta models (CompletionResponseChoice, CompletionResponseStreamChoice, ChatCompletionResponseChoice, DeltaMessage). Serialization logic was added to ensure hidden_states are omitted from the response if they are None.
• Hidden State Processing in Serving Endpoints: Implemented the core logic within serving_chat.py and serving_completions.py to capture and return hidden states. This includes passing the return_hidden_states flag to the internal request, collecting hidden states during generation, and including the last token's hidden states in both non-streaming and streaming responses.
• Utility Function for Hidden States: A new utility function, process_hidden_states_from_ret, was added to utils.py to standardize the extraction and processing of hidden states from the internal response (ret_item), specifically returning the hidden states of the last generated token.
• Comprehensive Unit Tests: Added new unit tests in test_serving_chat.py and test_serving_completions.py to validate the correct handling of the return_hidden_states request parameter, the proper processing and inclusion of hidden states in both non-streaming and streaming responses, and the functionality of the new utility function.

Changelog

• python/sglang/srt/entrypoints/openai/protocol.py
  • Added model_serializer import from pydantic.
  • Added return_hidden_states: bool = False to CompletionRequest (line 176) and ChatCompletionRequest (line 402).
  • Added hidden_states: Optional[object] = None to CompletionResponseChoice (line 212), CompletionResponseStreamChoice (line 237), ChatCompletionResponseChoice (line 464), and DeltaMessage (line 488).
  • Implemented @model_serializer(mode='wrap') methods for response choice/delta classes to remove hidden_states from the serialized output if their value is None.
• python/sglang/srt/entrypoints/openai/serving_chat.py
  • Imported process_hidden_states_from_ret from utils.py (line 33).
  • Passed request.return_hidden_states to the internal adapted_request (line 103).
  • Initialized hidden_states = {} for tracking in streaming generation (line 407).
  • Retrieved hidden_states from content['meta_info'] in the streaming loop (line 418).
  • Added logic to send a separate stream chunk containing the last token's hidden_states if request.return_hidden_states is true (lines 551-574).
  • Called process_hidden_states_from_ret to obtain hidden states for non-streaming responses (line 641).
  • Included hidden_states in the ChatCompletionResponseChoice construction (line 689).
• python/sglang/srt/entrypoints/openai/serving_completions.py
  • Imported process_hidden_states_from_ret from utils.py (line 23).
  • Passed request.return_hidden_states to the internal adapted_request (line 82).
  • Retrieved hidden_states from content['meta_info'] in the streaming loop (line 195).
  • Added logic to send a separate stream chunk containing the last token's hidden_states if request.return_hidden_states is true (lines 218-240).
  • Called process_hidden_states_from_ret to obtain hidden states for non-streaming responses (line 337).
  • Included hidden_states in the CompletionResponseChoice construction (line 351).
• python/sglang/srt/entrypoints/openai/usage_processor.py
  • Corrected the import path for UsageInfo from python.sglang.srt... to sglang.srt... (line 5).
• python/sglang/srt/entrypoints/openai/utils.py
  • Added imports for ChatCompletionRequest, CompletionRequest, and UsageInfo for type hinting (lines 7-11).
  • Introduced process_hidden_states_from_ret function to extract and process hidden states from a response item, returning the last token's hidden states (lines 216-241).
• test/srt/openai/test_serving_chat.py
  • Added self.mock_request setup in setUp (lines 65-66).
  • Added test_hidden_states_request_conversion to verify the return_hidden_states flag is correctly handled during request conversion (lines 227-270).
  • Added test_hidden_states_response_handling to test the process_hidden_states_from_ret utility function and non-streaming response processing (lines 271-330).
  • Added test_hidden_states_streaming to confirm hidden states are correctly included in streaming chat responses (lines 331-402).
• test/srt/openai/test_serving_completions.py
  • Added self.mock_request setup in setUp (line 35).
  • Added test_hidden_states_request_conversion to verify the return_hidden_states flag is correctly handled during request conversion (lines 101-119).
  • Added test_hidden_states_response_handling to test non-streaming completion responses with and without hidden states (lines 121-157).
  • Added test_hidden_states_streaming to confirm hidden states are correctly included in streaming completion responses (lines 159-215).

Activity

• On June 18, 2025, gemini-code-assist[bot] suggested improving type safety for hidden_states by using Optional[List[Any]] or Optional[List[float]] instead of Optional[object] in protocol.py.
• On June 18, 2025, gemini-code-assist[bot] recommended simplifying the condition for extracting last_token_hidden_states in both serving_chat.py and serving_completions.py by removing redundant checks.
• On June 19, 2025, CatherineSue validated the bot's suggestion for simplifying last_token_hidden_states and pointed out a redundant if choice_hidden_states: check in serving_chat.py.
• On June 19, 2025, CatherineSue inquired about the design choice for handling choice_hidden_states compared to an existing implementation in adapter.py.
• On June 19, 2025, key4ng (the author) confirmed that the redundant line 625 in serving_chat.py can be removed.
• On June 19, 2025, key4ng (the author) clarified that exclude_if_none in the Pydantic models would correctly handle cases where choice_hidden_states is None or empty, alleviating concerns about sending empty chunks.
• On June 19, 2025, CatherineSue noted the exclude_if_none pattern in protocol.py and suggested applying it consistently.
• On June 19, 2025, key4ng (the author) agreed on the need for consistency with exclude_if_none.
• On June 19, 2025, CatherineSue requested refinement of the request type hint in utils.py to remove List[*Request] as that type support has been removed.
• On June 20, 2025, CatherineSue approved the bot's suggestion to simplify the condition for hidden_states in utils.py.
• On June 20, 2025, key4ng (the author) initiated a /gemini summary command.