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

Co-authored-by: litao.dream <litao.dream@bytedance.com>

examples/runtime/hidden_states/hidden_states_engine.py

@@ -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()