From c1d529c19605cbf1f9be8db6d6d225b1465ea2e0 Mon Sep 17 00:00:00 2001 From: Aurick Qiao Date: Tue, 3 Feb 2026 18:33:27 -0800 Subject: [PATCH] Fix Session for multimodal and expose it through Engine (#18152) --- python/sglang/srt/entrypoints/engine.py | 37 ++++++++++++++++++++ python/sglang/srt/managers/schedule_batch.py | 1 - python/sglang/srt/managers/scheduler.py | 17 +++++++++ 3 files changed, 54 insertions(+), 1 deletion(-) diff --git a/python/sglang/srt/entrypoints/engine.py b/python/sglang/srt/entrypoints/engine.py index 8d7ab8716..1ad96d6cf 100644 --- a/python/sglang/srt/entrypoints/engine.py +++ b/python/sglang/srt/entrypoints/engine.py @@ -42,6 +42,7 @@ from sglang.srt.managers.data_parallel_controller import ( ) from sglang.srt.managers.detokenizer_manager import run_detokenizer_process from sglang.srt.managers.io_struct import ( + CloseSessionReqInput, DestroyWeightsUpdateGroupReqInput, EmbeddingReqInput, GenerateReqInput, @@ -50,6 +51,7 @@ from sglang.srt.managers.io_struct import ( LoadLoRAAdapterFromTensorsReqInput, LoadLoRAAdapterReqInput, MultimodalDataInputFormat, + OpenSessionReqInput, ReleaseMemoryOccupationReqInput, ResumeMemoryOccupationReqInput, RpcReqInput, @@ -233,6 +235,7 @@ class Engine(EngineBase): data_parallel_rank: Optional[int] = None, external_trace_header: Optional[Dict] = None, rid: Optional[Union[List[str], str]] = None, + session_params: Optional[Dict] = None, ) -> Union[Dict, Iterator[Dict]]: """ The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`. @@ -270,6 +273,7 @@ class Engine(EngineBase): data_parallel_rank=data_parallel_rank, external_trace_header=external_trace_header, rid=rid, + session_params=session_params, ) generator = self.tokenizer_manager.generate_request(obj, None) @@ -320,6 +324,7 @@ class Engine(EngineBase): data_parallel_rank: Optional[int] = None, external_trace_header: Optional[Dict] = None, rid: Optional[Union[List[str], str]] = None, + session_params: Optional[Dict] = None, ) -> Union[Dict, AsyncIterator[Dict]]: """ The arguments of this function is the same as `sglang/srt/managers/io_struct.py::GenerateReqInput`. @@ -358,6 +363,7 @@ class Engine(EngineBase): data_parallel_rank=data_parallel_rank, external_trace_header=external_trace_header, rid=rid, + session_params=session_params, ) generator = self.tokenizer_manager.generate_request(obj, None) @@ -448,6 +454,37 @@ class Engine(EngineBase): def flush_cache(self): return self.loop.run_until_complete(self.tokenizer_manager.flush_cache()) + def open_session( + self, + capacity_of_str_len: int, + session_id: Optional[str] = None, + ) -> str: + """Open a session for multi-turn conversation with shared context. + + Args: + capacity_of_str_len: Maximum string length capacity for the session. + session_id: Optional session ID. If not provided, a UUID will be generated. + + Returns: + The session ID (either the provided one or a newly generated UUID). + """ + obj = OpenSessionReqInput( + capacity_of_str_len=capacity_of_str_len, + session_id=session_id, + ) + return self.loop.run_until_complete( + self.tokenizer_manager.open_session(obj, None) + ) + + def close_session(self, session_id: str) -> None: + """Close a session and release its resources. + + Args: + session_id: The session ID to close. + """ + obj = CloseSessionReqInput(session_id=session_id) + self.loop.run_until_complete(self.tokenizer_manager.close_session(obj, None)) + def start_profile(self, **kwargs): self.loop.run_until_complete(self.tokenizer_manager.start_profile(**kwargs)) diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 5768b28d2..3181c9749 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/python/sglang/srt/managers/schedule_batch.py @@ -1107,7 +1107,6 @@ class Req: self.last_node = None self.swa_uuid_for_lock = None self.extend_input_len = 0 - self.customized_info = None self.is_retracted = True self.retracted_stain = True self.input_token_logprobs = None diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index e59db607e..c7c704c1b 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1415,6 +1415,10 @@ class Scheduler( for req in batch.reqs: if not req.finished() or not (mm_inputs := req.multimodal_inputs): continue + # For session requests, keep mm_inputs for the next request + if req.session_id: + continue + # For non-session requests, clear features and mm_inputs for item in mm_inputs.mm_items: item.feature = None req.multimodal_inputs = None @@ -1508,6 +1512,19 @@ class Scheduler( if recv_req.mm_inputs is not None: image_inputs = self._get_multimodal_inputs(recv_req.mm_inputs) + # For session requests, adjust mm_inputs offsets by the prefix length. + # Session.create_req prepends previous context to origin_input_ids, + # so offsets from the new prompt need to be shifted. + if len(recv_req.input_ids) < len(req.origin_input_ids): + assert recv_req.session_params.id in self.sessions + prefix_len = len(req.origin_input_ids) - len(recv_req.input_ids) + for mm_item in image_inputs.mm_items: + if mm_item.offsets: + mm_item.offsets = [ + (start + prefix_len, end + prefix_len) + for start, end in mm_item.offsets + ] + # The following steps are already fast, execute locally on each rank. # Expand a single image token into multiple dummy tokens for receiving image embeddings req.origin_input_ids = self.pad_input_ids_func(