From 20beb3702b178de2ccd4240f949b76f9cf8a3eb8 Mon Sep 17 00:00:00 2001
From: eigen <52445717+yyihuang@users.noreply.github.com>
Date: Wed, 25 Jun 2025 05:10:09 -0400
Subject: [PATCH] feat: add return hidden_states at async generation (#7507)

---
 python/sglang/srt/disaggregation/mooncake/conn.py | 2 +-
 python/sglang/srt/entrypoints/engine.py           | 2 ++
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/python/sglang/srt/disaggregation/mooncake/conn.py b/python/sglang/srt/disaggregation/mooncake/conn.py
index 1d6219089..3c05de5e0 100644
--- a/python/sglang/srt/disaggregation/mooncake/conn.py
+++ b/python/sglang/srt/disaggregation/mooncake/conn.py
@@ -422,7 +422,7 @@ class MooncakeKVManager(BaseKVManager):
                     src_addr_list.append(src_slice_addr)
                     dst_addr_list.append(dst_slice_addr)
                     length_list.append(slice_lens_per_page)
-                
+
                     logger.debug(
                         f"SYNC: sid={mooncake_session_id}, "
                         f"src={src_slice_addr}, dst={dst_slice_addr}, len={slice_lens_per_page}"
diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py
index 46dbc31fc..0f75b2380 100644
--- a/python/sglang/srt/entrypoints/engine.py
+++ b/python/sglang/srt/entrypoints/engine.py
@@ -242,6 +242,7 @@ class Engine(EngineBase):
         token_ids_logprob: Optional[Union[List[List[int]], List[int]]] = None,
         lora_path: Optional[List[Optional[str]]] = None,
         custom_logit_processor: Optional[Union[List[str], str]] = None,
+        return_hidden_states: bool = False,
         stream: bool = False,
         bootstrap_host: Optional[Union[List[str], str]] = None,
         bootstrap_port: Optional[Union[List[int], int]] = None,
@@ -274,6 +275,7 @@ class Engine(EngineBase):
             top_logprobs_num=top_logprobs_num,
             token_ids_logprob=token_ids_logprob,
             lora_path=lora_path,
+            return_hidden_states=return_hidden_states,
             stream=stream,
             custom_logit_processor=custom_logit_processor,
             bootstrap_host=bootstrap_host,