minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 10 Packages Projects Releases Wiki Activity

[Feature] Support return_hidden_states="last" (#30177)

Browse Source 
Co-authored-by: litao.dream <litao.dream@bytedance.com>
This commit is contained in:
Tao Li
2026-08-02 15:09:33 +08:00
committed by GitHub
co-authored by litao.dream
parent 1685d29f21
commit a0b7bcf592
30 changed files with 1096 additions and 268 deletions
Download Patch File Download Diff File Expand all files Collapse all files
examples/runtime/hidden_states/hidden_states_engine.py
+10 -18
View File
@@ -1,10 +1,9 @@
"""
Usage:
python hidden_states.py
python hidden_states_engine.py
Note that each time you change the `return_hidden_states` parameter,
the cuda graph will be recaptured, which might lead to a performance hit.
So avoid getting hidden states and completions alternately.
CUDA graphs use the configured maximum hidden-state mode. Requests may select
that mode or a weaker one without triggering mode-dependent recapture.
"""
import torch
@@ -22,7 +21,7 @@ def main():
# Create an LLM.
llm = sgl.Engine(
model_path="Alibaba-NLP/gte-Qwen2-1.5B-instruct",
enable_return_hidden_states=True,
return_hidden_states_mode="last",
)
sampling_params = {
@@ -32,16 +31,15 @@ def main():
}
outputs = llm.generate(
prompts, sampling_params=sampling_params, return_hidden_states=True
prompts, sampling_params=sampling_params, return_hidden_states="last"
)
llm.shutdown()
for prompt, output in zip(prompts, outputs):
for i in range(len(output["meta_info"]["hidden_states"])):
output["meta_info"]["hidden_states"][i] = torch.tensor(
output["meta_info"]["hidden_states"][i], dtype=torch.bfloat16
)
hidden_state = torch.tensor(
output["meta_info"]["hidden_states"], dtype=torch.bfloat16
)
print("===============================")
print(
f"Prompt: {prompt}\n"
@@ -49,14 +47,8 @@ def main():
f"Prompt_Tokens: {output['meta_info']['prompt_tokens']}\t"
f"Completion_tokens: {output['meta_info']['completion_tokens']}"
)
print("Hidden states: ")
hidden_states = torch.cat(
[
i.unsqueeze(0) if len(i.shape) == 1 else i
for i in output["meta_info"]["hidden_states"]
]
)
print(hidden_states)
print("Last hidden state: ")
print(hidden_state)
print()
examples/runtime/hidden_states/hidden_states_server.py
+11 -17
View File
@@ -3,9 +3,8 @@ Usage:
python hidden_states_server.py
Note that each time you change the `return_hidden_states` parameter,
the cuda graph will be recaptured, which might lead to a performance hit.
So avoid getting hidden states and completions alternately.
CUDA graphs use the configured maximum hidden-state mode. Requests may select
that mode or a weaker one without triggering mode-dependent recapture.
"""
import requests
@@ -23,7 +22,9 @@ else:
def main():
# Launch the server
server_process, port = launch_server_cmd(
"python -m sglang.launch_server --model-path Alibaba-NLP/gte-Qwen2-1.5B-instruct --enable-return-hidden-states --host 0.0.0.0"
"python -m sglang.launch_server --model-path "
"Alibaba-NLP/gte-Qwen2-1.5B-instruct "
"--return-hidden-states-mode last --host 0.0.0.0"
)
wait_for_server(f"http://localhost:{port}", process=server_process)
@@ -43,7 +44,7 @@ def main():
json_data = {
"text": prompts,
"sampling_params": sampling_params,
"return_hidden_states": True,
"return_hidden_states": "last",
}
response = requests.post(
@@ -55,10 +56,9 @@ def main():
outputs = response.json()
for prompt, output in zip(prompts, outputs):
for i in range(len(output["meta_info"]["hidden_states"])):
output["meta_info"]["hidden_states"][i] = torch.tensor(
output["meta_info"]["hidden_states"][i], dtype=torch.bfloat16
)
hidden_state = torch.tensor(
output["meta_info"]["hidden_states"], dtype=torch.bfloat16
)
print("===============================")
print(
f"Prompt: {prompt}\n"
@@ -66,14 +66,8 @@ def main():
f"Prompt_Tokens: {output['meta_info']['prompt_tokens']}\t"
f"Completion_tokens: {output['meta_info']['completion_tokens']}"
)
print("Hidden states: ")
hidden_states = torch.cat(
[
i.unsqueeze(0) if len(i.shape) == 1 else i
for i in output["meta_info"]["hidden_states"]
]
)
print(hidden_states)
print("Last hidden state: ")
print(hidden_state)
print()