[Feat] return hidden states

[Jackmin801]

Motivation

This PR intends to add the return_hidden_states argument to ServerArgs which makes the results contain the last layer hidden states in output["meta_info"]["hidden_states"].
These hidden states are useful for example for verifying computations. (e.g. https://arxiv.org/abs/2501.16007)

Modifications

• Add return_hidden_states to ServerArgs
• Changed the logic to determine capture_hidden_mode to accomodate return_hidden_states
• Modify scheduler process_batch_results to save the hidden state to the Req
• Add return_hidden_states and hidden_states to necessary dataclasses

Script used to test changes

# launch the offline engine
import asyncio
from transformers import AutoTokenizer
import sglang as sgl

def main():
    MODEL_NAME = "meta-llama/Meta-Llama-3.1-8B-Instruct"
    llm = sgl.Engine(
        model_path=MODEL_NAME,
        skip_tokenizer_init=True,
        disable_cuda_graph=False,
        return_hidden_states=False,
    )
    tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)

    prompts = [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
        "The future of AI is",
    ]

    sampling_params = {"temperature": 0.8, "top_p": 0.95, "max_new_tokens": 10}

    input_ids = tokenizer(prompts).input_ids
    #outputs = llm.generate(input_ids=input_ids, sampling_params=sampling_params)
    outputs = llm.generate(prompts, sampling_params=sampling_params)
    for input_id, output in zip(input_ids, outputs):
        print("===============================")
        print(input_id)
        print(output)
        print()
        if "token_ids" in output:
            print(input_id, output["token_ids"], len(input_id), len(output["token_ids"]))
        else:
            print(output['text'])
        if "hidden_states" in output["meta_info"]:
            print(
                [i.shape for i in output["meta_info"]["hidden_states"]],
                len(output["meta_info"]["hidden_states"]),
            )

if __name__ == "__main__":
    main()

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.

[zhaochenyang20]

This is good to see. But could change our documents to demonstrate the usage and add unit tests to your feature?

[zhaochenyang20]

Thanks. I will try to get some one familiar with hidden state to help.

[zhaochenyang20]

https://github.com/sgl-project/sglang/blob/main/test/srt/models/test_generation_models.py

You can check this, may gonna help.

[zhaochenyang20]

Great! Also, could you add the API also to the server, not only the engine. Like how we do for update_weights_from_dist. You can use Engine API and Server / HTTPS API.

[zhaochenyang20]

Thanks. I think we should not change the docs of offline API. Instead, we should change this.

https://docs.sglang.ai/backend/native_api.html

Also, I think that the best to do this is not add an serving arguement, but rather make a new native API instead. Just like:

sglang/python/sglang/srt/entrypoints/http_server.py

Line 291 in f90db8b

@app.post("/update_weights_from_distributed")

And this:

sglang/python/sglang/srt/entrypoints/engine.py

Line 249 in f90db8b

def update_weights_from_distributed(self, name: str, dtype, shape):

This could be much easier to use and do not need to launch a specific engine, which cost a lot of time in the docs CI.

[zhaochenyang20]

Also, update the beginning of the native API docs.

https://docs.sglang.ai/backend/native_api.html

Apart from the OpenAI compatible APIs, the SGLang Runtime also provides its native server APIs. We introduce these following APIs:

/generate (text generation model)

/get_model_info

/get_server_info

/health

/health_generate

/flush_cache

/update_weights

/encode(embedding model)

/classify(reward model)

We mainly use requests to test these APIs in the following examples. You can also use curl.

[zhaochenyang20]

You can add one seperate test file as test_hidden_state.py, but add it in https://github.com/sgl-project/sglang/blob/main/test/srt/run_suite.py

Just like test_update_weights_from_disk in the run_suite.

[zhaochenyang20]

The test looks good.

[Jackmin801]

@zhaochenyang20 What I meant is here. It seems like it is necessary for the capture_hidden_mode to be known at engine init time. Otherwise, the decode cuda graph will not contain the return hidden state logic and this cant be changed by sampling args.

[zhaochenyang20]

https://github.com/sgl-project/sglang/actions/runs/13235087533/job/36938449704?pr=3364

This needs to update the time out of the CI @Jackmin801