sglang 0.4.7.post1__py3-none-any.whl → 0.4.8.post1__py3-none-any.whl

sglang/srt/disaggregation/nixl/conn.py

@@ -31,23 +31,19 @@ from sglang.srt.utils import get_local_ip_by_remote
 
 logger = logging.getLogger(__name__)
 
-NixlEngineInfo: TypeAlias = Dict[str, Union[str, int]]
-
 GUARD = "NixlMsgGuard".encode("ascii")
 
 
 @dataclasses.dataclass
 class TransferInfo:
+    """Contains indices for a transfer, sent by KVReceiver. Received by prefill bootstrap thread."""
+
     room: int
     endpoint: str
     dst_port: int
-    agent_metadata: bytes
     agent_name: str
-    dst_kv_ptrs: list[int]
     dst_kv_indices: npt.NDArray[np.int32]
-    dst_aux_ptrs: list[int]
     dst_aux_index: int
-    dst_gpu_id: int
     required_dst_info_num: int
 
     def is_dummy(self):
@@ -59,14 +55,37 @@ class TransferInfo:
             room=int(msg[0].decode("ascii")),
             endpoint=msg[1].decode("ascii"),
             dst_port=int(msg[2].decode("ascii")),
-            agent_metadata=msg[3],
-            agent_name=msg[4].decode("ascii"),
+            agent_name=msg[3].decode("ascii"),
+            dst_kv_indices=np.frombuffer(msg[4], dtype=np.int32),
+            dst_aux_index=int(msg[5].decode("ascii")),
+            required_dst_info_num=int(msg[6].decode("ascii")),
+        )
+
+
+@dataclasses.dataclass
+class KVArgsRegisterInfo:
+    """Contains base pointers and other info which only needs to be sent once by KVReceiver. Received by prefill bootstrap thread."""
+
+    room: str
+    endpoint: str
+    dst_port: int
+    agent_name: str
+    agent_metadata: bytes
+    dst_kv_ptrs: list[int]
+    dst_aux_ptrs: list[int]
+    gpu_id: int
+
+    @classmethod
+    def from_zmq(cls, msg: List[bytes]):
+        return cls(
+            room=str(msg[0].decode("ascii")),
+            endpoint=msg[1].decode("ascii"),
+            dst_port=int(msg[2].decode("ascii")),
+            agent_name=msg[3].decode("ascii"),
+            agent_metadata=msg[4],
             dst_kv_ptrs=list(struct.unpack(f"{len(msg[5])//8}Q", msg[5])),
-            dst_kv_indices=np.frombuffer(msg[6], dtype=np.int32),
-            dst_aux_ptrs=list(struct.unpack(f"{len(msg[7])//8}Q", msg[7])),
-            dst_aux_index=int(msg[8].decode("ascii")),
-            dst_gpu_id=int(msg[9].decode("ascii")),
-            required_dst_info_num=int(msg[10].decode("ascii")),
+            dst_aux_ptrs=list(struct.unpack(f"{len(msg[6])//8}Q", msg[6])),
+            gpu_id=int(msg[7].decode("ascii")),
         )
 
 
@@ -109,9 +128,9 @@ class NixlKVManager(CommonKVManager):
         self.register_buffer_to_engine()
 
         if self.disaggregation_mode == DisaggregationMode.PREFILL:
-            self.request_status = {}
-            self.transfer_infos: Dict[int, TransferInfo] = {}
-            self.peer_names: Dict[str, str] = {}
+            self.request_status: Dict[int, KVPoll] = {}
+            self.transfer_infos: Dict[int, Dict[str, TransferInfo]] = {}
+            self.decode_kv_args_table: Dict[str, KVArgsRegisterInfo] = {}
             self._start_bootstrap_thread()
         elif self.disaggregation_mode == DisaggregationMode.DECODE:
             self.transfer_statuses: Dict[int, TransferStatus] = defaultdict(
@@ -154,10 +173,13 @@ class NixlKVManager(CommonKVManager):
         if not self.aux_descs:
             raise Exception("NIXL memory registration failed for aux tensors")
 
-    def _add_remote(self, agent_name: str, agent_metadata: bytes):
-        if agent_name not in self.peer_names:
-            self.peer_names[agent_name] = self.agent.add_remote_agent(agent_metadata)
-        return self.peer_names[agent_name]
+    def _add_remote_peer(self, decode_kv_args: KVArgsRegisterInfo):
+        agent_name = decode_kv_args.agent_name
+        if agent_name in self.decode_kv_args_table:
+            logger.info(f"Peer {agent_name} was already registered, ignoring.")
+            return
+        self.decode_kv_args_table[agent_name] = decode_kv_args
+        self.agent.add_remote_agent(decode_kv_args.agent_metadata)
 
     def send_kvcache(
         self,
@@ -262,17 +284,17 @@ class NixlKVManager(CommonKVManager):
             if req.is_dummy():
                 continue
 
-            peer_name = self._add_remote(req.agent_name, req.agent_metadata)
             chunked_dst_kv_indice = req.dst_kv_indices[index_slice]
             assert len(chunked_dst_kv_indice) == len(kv_indices)
+            assert req.agent_name in self.decode_kv_args_table
 
             notif = "_".join([str(req.room), "kv", str(chunk_id), str(int(is_last))])
             kv_xfer_handle = self.send_kvcache(
-                peer_name,
+                req.agent_name,
                 kv_indices,
-                req.dst_kv_ptrs,
+                self.decode_kv_args_table[req.agent_name].dst_kv_ptrs,
                 chunked_dst_kv_indice,
-                req.dst_gpu_id,
+                self.decode_kv_args_table[req.agent_name].gpu_id,
                 notif,
             )
             handles.append(kv_xfer_handle)
@@ -280,13 +302,15 @@ class NixlKVManager(CommonKVManager):
             if is_last:
                 assert aux_index is not None
                 aux_xfer_handle = self.send_aux(
-                    peer_name,
+                    req.agent_name,
                     aux_index,
-                    req.dst_aux_ptrs,
+                    self.decode_kv_args_table[req.agent_name].dst_aux_ptrs,
                     req.dst_aux_index,
                     str(req.room) + "_aux",
                 )
                 handles.append(aux_xfer_handle)
+        if is_last:
+            del self.transfer_infos[bootstrap_room]
         return handles
 
     def update_transfer_status(self):
@@ -328,16 +352,23 @@ class NixlKVManager(CommonKVManager):
                 ), f"First message should be {GUARD}. Foreign traffic?"
                 waiting_req_bytes = waiting_req_bytes[1:]
                 room = waiting_req_bytes[0].decode("ascii")
-
-                required_dst_info_num = int(waiting_req_bytes[10].decode("ascii"))
+                agent_name = waiting_req_bytes[3].decode("ascii")
+                if room == "None":
+                    # Register new peer and save KV base pointers.
+                    self._add_remote_peer(
+                        KVArgsRegisterInfo.from_zmq(waiting_req_bytes)
+                    )
+                    logger.debug(f"Register KVArgs from {agent_name} successfully")
+                    continue
                 room = int(room)
-                agent_name = waiting_req_bytes[4].decode("ascii")
                 if room not in self.transfer_infos:
                     self.transfer_infos[room] = {}
                 self.transfer_infos[room][agent_name] = TransferInfo.from_zmq(
                     waiting_req_bytes
                 )
-
+                required_dst_info_num = self.transfer_infos[room][
+                    agent_name
+                ].required_dst_info_num
                 logger.debug(f"got info {room=} {agent_name=} {required_dst_info_num=}")
                 if len(self.transfer_infos[room]) == required_dst_info_num:
                     logger.debug(f"{room=} is bootstrapped")
@@ -391,6 +422,7 @@ class NixlKVSender(BaseKVSender):
         self.chunk_id += 1
         if is_last:
             self.has_sent = True
+            del self.kv_mgr.request_status[self.bootstrap_room]
 
     def poll(self) -> KVPoll:
         if not self.has_sent:
@@ -415,6 +447,7 @@ class NixlKVReceiver(CommonKVReceiver):
         data_parallel_rank: Optional[int] = None,
     ):
         self.started_transfer = False
+        self.conclude_state = None
         super().__init__(mgr, bootstrap_addr, bootstrap_room, data_parallel_rank)
 
     def init(self, kv_indices: npt.NDArray[np.int32], aux_index: Optional[int] = None):
@@ -426,17 +459,8 @@ class NixlKVReceiver(CommonKVReceiver):
                 f"Fetched bootstrap info: {bootstrap_info} for engine rank: {self.kv_mgr.kv_args.engine_rank}"
             )
             is_dummy = bootstrap_info["is_dummy"]
-
-            # TODO: send_kv_args earlier
-            packed_kv_data_ptrs = b"".join(
-                struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs
-            )
-            packed_aux_data_ptrs = b"".join(
-                struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs
-            )
-
             logger.debug(
-                f"Sending to {self.prefill_server_url} with bootstrap room {self.bootstrap_room}"
+                f"Sending to {self.prefill_server_url} with bootstrap room {self.bootstrap_room} {is_dummy=}"
             )
             sock, lock = self._connect("tcp://" + self.prefill_server_url)
             with lock:
@@ -446,13 +470,9 @@ class NixlKVReceiver(CommonKVReceiver):
                         str(self.bootstrap_room).encode("ascii"),
                         get_local_ip_by_remote().encode("ascii"),
                         str(self.kv_mgr.rank_port).encode("ascii"),
-                        self.kv_mgr.agent.get_agent_metadata(),
                         self.kv_mgr.agent.name.encode("ascii"),
-                        packed_kv_data_ptrs,
                         kv_indices.tobytes() if not is_dummy else b"",
-                        packed_aux_data_ptrs,
                         str(aux_index).encode("ascii"),
-                        str(self.kv_mgr.kv_args.gpu_id).encode("ascii"),
                         str(self.required_dst_info_num).encode("ascii"),
                     ]
                 )
@@ -460,17 +480,45 @@ class NixlKVReceiver(CommonKVReceiver):
         self.started_transfer = True
 
     def poll(self) -> KVPoll:
+        if self.conclude_state is not None:
+            return self.conclude_state
         if not self.started_transfer:
             return KVPoll.WaitingForInput  # type: ignore
 
         self.kv_mgr.update_transfer_status()
-
         if self.kv_mgr.check_transfer_done(self.bootstrap_room):  # type: ignore
+            self.conclude_state = KVPoll.Success
+            del self.kv_mgr.transfer_statuses[self.bootstrap_room]
             return KVPoll.Success  # type: ignore
         return KVPoll.WaitingForInput  # type: ignore
 
     def _register_kv_args(self):
-        pass
+        for bootstrap_info in self.bootstrap_infos:
+            self.prefill_server_url = (
+                f"{bootstrap_info['rank_ip']}:{bootstrap_info['rank_port']}"
+            )
+            packed_kv_data_ptrs = b"".join(
+                struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.kv_data_ptrs
+            )
+            packed_aux_data_ptrs = b"".join(
+                struct.pack("Q", ptr) for ptr in self.kv_mgr.kv_args.aux_data_ptrs
+            )
+
+            sock, lock = self._connect("tcp://" + self.prefill_server_url)
+            with lock:
+                sock.send_multipart(
+                    [
+                        GUARD,
+                        "None".encode("ascii"),
+                        get_local_ip_by_remote().encode("ascii"),
+                        str(self.kv_mgr.rank_port).encode("ascii"),
+                        self.kv_mgr.agent.name.encode("ascii"),
+                        self.kv_mgr.agent.get_agent_metadata(),
+                        packed_kv_data_ptrs,
+                        packed_aux_data_ptrs,
+                        str(self.kv_mgr.kv_args.gpu_id).encode("ascii"),
+                    ]
+                )
 
     def failure_exception(self):
         raise Exception("Fake KVReceiver Exception")

sglang/srt/disaggregation/prefill.py

@@ -25,7 +25,6 @@ from collections import deque
 from http import HTTPStatus
 from typing import TYPE_CHECKING, List, Optional
 
-import numpy as np
 import torch
 
 from sglang.srt.disaggregation.base import BaseKVManager, KVPoll
@@ -45,6 +44,7 @@ from sglang.srt.disaggregation.utils import (
 )
 from sglang.srt.managers.schedule_batch import FINISH_LENGTH, Req, ScheduleBatch
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
+from sglang.srt.utils import require_mlp_sync
 
 if TYPE_CHECKING:
     from torch.distributed import ProcessGroup
@@ -93,8 +93,6 @@ class PrefillBootstrapQueue:
         self.gpu_id = gpu_id
         self.bootstrap_port = bootstrap_port
         self.queue: List[Req] = []
-        self.pp_rank = pp_rank
-        self.pp_size = pp_size
         self.gloo_group = gloo_group
         self.max_total_num_tokens = max_total_num_tokens
         self.scheduler = scheduler
@@ -124,6 +122,9 @@ class PrefillBootstrapQueue:
         kv_args.kv_data_ptrs = kv_data_ptrs
         kv_args.kv_data_lens = kv_data_lens
         kv_args.kv_item_lens = kv_item_lens
+        if not self.is_mla_backend:
+            kv_args.kv_head_num = self.token_to_kv_pool.head_num
+        kv_args.page_size = self.token_to_kv_pool.page_size
 
         kv_args.aux_data_ptrs, kv_args.aux_data_lens, kv_args.aux_item_lens = (
             self.metadata_buffers.get_buf_infos()
@@ -274,12 +275,8 @@ class SchedulerDisaggregationPrefillMixin:
             self.process_prefill_chunk()
             batch = self.get_new_batch_prefill()
 
-            # Handle DP attention
-            if (
-                self.server_args.enable_dp_attention
-                or self.server_args.enable_sp_layernorm
-            ):
-                batch, _ = self.prepare_dp_attn_batch(batch)
+            if require_mlp_sync(self.server_args):
+                batch, _ = self.prepare_mlp_sync_batch(batch)
             self.cur_batch = batch
 
             if batch:
@@ -312,12 +309,8 @@ class SchedulerDisaggregationPrefillMixin:
             self.process_prefill_chunk()
             batch = self.get_new_batch_prefill()
 
-            # Handle DP attention
-            if (
-                self.server_args.enable_dp_attention
-                or self.server_args.enable_sp_layernorm
-            ):
-                batch, _ = self.prepare_dp_attn_batch(batch)
+            if require_mlp_sync(self.server_args):
+                batch, _ = self.prepare_mlp_sync_batch(batch)
             self.cur_batch = batch
             if batch:
                 result = self.run_batch(batch)
@@ -393,6 +386,8 @@ class SchedulerDisaggregationPrefillMixin:
                     logits_output.input_token_logprobs = tuple(
                         logits_output.input_token_logprobs.tolist()
                     )
+
+        hidden_state_offset = 0
         for i, (req, next_token_id) in enumerate(
             zip(batch.reqs, next_token_ids, strict=True)
         ):
@@ -402,6 +397,16 @@ class SchedulerDisaggregationPrefillMixin:
                 req.output_ids.append(next_token_id)
                 self.tree_cache.cache_unfinished_req(req)  # update the tree and lock
                 self.disagg_prefill_inflight_queue.append(req)
+                if logits_output.hidden_states is not None:
+                    last_hidden_index = (
+                        hidden_state_offset + extend_input_len_per_req[i] - 1
+                    )
+                    req.hidden_states_tensor = (
+                        logits_output.hidden_states[last_hidden_index].cpu().clone()
+                    )
+                    hidden_state_offset += extend_input_len_per_req[i]
+                else:
+                    req.hidden_states_tensor = None
                 if req.return_logprob:
                     assert extend_logprob_start_len_per_req is not None
                     assert extend_input_len_per_req is not None

sglang/srt/disaggregation/utils.py

@@ -6,6 +6,7 @@ import random
 import threading
 import warnings
 from collections import deque
+from contextlib import nullcontext
 from enum import Enum
 from typing import TYPE_CHECKING, List, Optional
 
@@ -84,24 +85,43 @@ class ReqToMetadataIdxAllocator:
 
 
 class MetadataBuffers:
-    def __init__(self, size: int, max_top_logprobs_num: int = 128):
-        # TODO: abort top_logprobs_num > 128 in PD
-
-        # We transfer the metadata of first output token to decode
-        # The minimal size for RDMA is 64Bytes, so we pad it to > 64Bytes
-        self.output_ids = torch.zeros((size, 16), dtype=torch.int32, device="cpu")
-        self.output_token_logprobs_val = torch.zeros(
-            (size, 16), dtype=torch.float32, device="cpu"
-        )
-        self.output_token_logprobs_idx = torch.zeros(
-            (size, 16), dtype=torch.int32, device="cpu"
-        )
-        self.output_top_logprobs_val = torch.zeros(
-            (size, max_top_logprobs_num), dtype=torch.float32, device="cpu"
-        )
-        self.output_top_logprobs_idx = torch.zeros(
-            (size, max_top_logprobs_num), dtype=torch.int32, device="cpu"
-        )
+    def __init__(
+        self,
+        size: int,
+        hidden_size: int,
+        dtype: torch.dtype,
+        max_top_logprobs_num: int = 128,
+        custom_mem_pool: torch.cuda.MemPool = None,
+    ):
+        self.custom_mem_pool = custom_mem_pool
+        device = "cuda" if self.custom_mem_pool else "cpu"
+
+        with (
+            torch.cuda.use_mem_pool(self.custom_mem_pool)
+            if self.custom_mem_pool
+            else nullcontext()
+        ):
+            # TODO: abort top_logprobs_num > 128 in PD
+
+            # We transfer the metadata of first output token to decode
+            # The minimal size for RDMA is 64Bytes, so we pad it to > 64Bytes
+            self.output_ids = torch.zeros((size, 16), dtype=torch.int32, device=device)
+
+            self.output_token_logprobs_val = torch.zeros(
+                (size, 16), dtype=torch.float32, device=device
+            )
+            self.output_token_logprobs_idx = torch.zeros(
+                (size, 16), dtype=torch.int32, device=device
+            )
+            self.output_top_logprobs_val = torch.zeros(
+                (size, max_top_logprobs_num), dtype=torch.float32, device=device
+            )
+            self.output_top_logprobs_idx = torch.zeros(
+                (size, max_top_logprobs_num), dtype=torch.int32, device=device
+            )
+            self.output_hidden_states = torch.zeros(
+                (size, hidden_size), dtype=dtype, device=device
+            )
 
     def get_buf_infos(self):
         ptrs = [
@@ -110,6 +130,7 @@ class MetadataBuffers:
             self.output_token_logprobs_idx.data_ptr(),
             self.output_top_logprobs_val.data_ptr(),
             self.output_top_logprobs_idx.data_ptr(),
+            self.output_hidden_states.data_ptr(),
         ]
         data_lens = [
             self.output_ids.nbytes,
@@ -117,6 +138,7 @@ class MetadataBuffers:
             self.output_token_logprobs_idx.nbytes,
             self.output_top_logprobs_val.nbytes,
             self.output_top_logprobs_idx.nbytes,
+            self.output_hidden_states.nbytes,
         ]
         item_lens = [
             self.output_ids[0].nbytes,
@@ -124,6 +146,7 @@ class MetadataBuffers:
             self.output_token_logprobs_idx[0].nbytes,
             self.output_top_logprobs_val[0].nbytes,
             self.output_top_logprobs_idx[0].nbytes,
+            self.output_hidden_states[0].nbytes,
         ]
         return ptrs, data_lens, item_lens
 
@@ -134,6 +157,7 @@ class MetadataBuffers:
             self.output_token_logprobs_idx[idx],
             self.output_top_logprobs_val[idx],
             self.output_top_logprobs_idx[idx],
+            self.output_hidden_states[idx],
         )
 
     def set_buf(self, req: Req):
@@ -161,6 +185,11 @@ class MetadataBuffers:
                 ] = torch.tensor(
                     req.output_top_logprobs_idx[0], dtype=torch.int32, device="cpu"
                 )
+        # for PD + spec decode
+        if req.hidden_states_tensor is not None:
+            self.output_hidden_states[req.metadata_buffer_index].copy_(
+                req.hidden_states_tensor
+            )
 
 
 #########################

sglang/srt/distributed/parallel_state.py

@@ -523,17 +523,25 @@ class GroupCoordinator:
         self,
         input_: torch.Tensor,
         dim: int = -1,
-        tensor_list: List[torch.Tensor] = None,
+        output_tensor_list: Optional[List[torch.Tensor]] = None,
     ) -> torch.Tensor:
         world_size = self.world_size
         # Bypass the function if we are using only 1 GPU.
         if world_size == 1:
-            return input_
+            if output_tensor_list is not None:
+                logger.warning(
+                    "Performing in-place all-gather with a group size of 1. "
+                    "This may be unnecessary; consider bypassing it for better efficiency."
+                )
+                output_tensor_list[0].copy_(input_)
+                return None
+            else:
+                return input_
 
-        if tensor_list is not None:
+        if output_tensor_list is not None:
             # TODO(ch-wan): support other backends
             return torch.distributed.all_gather(
-                tensor_list, input_, group=self.device_group
+                output_tensor_list, input_, group=self.device_group
             )
 
         assert (

sglang/srt/entrypoints/engine.py

@@ -37,7 +37,6 @@ setattr(threading, "_register_atexit", lambda *args, **kwargs: None)
 import torch
 import uvloop
 
-from sglang.srt.code_completion_parser import load_completion_template_for_openai_api
 from sglang.srt.entrypoints.EngineBase import EngineBase
 from sglang.srt.managers.data_parallel_controller import (
     run_data_parallel_controller_process,
@@ -58,11 +57,8 @@ from sglang.srt.managers.io_struct import (
     UpdateWeightsFromTensorReqInput,
 )
 from sglang.srt.managers.scheduler import run_scheduler_process
+from sglang.srt.managers.template_manager import TemplateManager
 from sglang.srt.managers.tokenizer_manager import TokenizerManager
-from sglang.srt.openai_api.adapter import (
-    guess_chat_template_name_from_model_path,
-    load_chat_template_for_openai_api,
-)
 from sglang.srt.server_args import PortArgs, ServerArgs
 from sglang.srt.torch_memory_saver_adapter import TorchMemorySaverAdapter
 from sglang.srt.utils import (
@@ -119,21 +115,22 @@ class Engine(EngineBase):
         atexit.register(self.shutdown)
 
         # Allocate ports for inter-process communications
-        port_args = PortArgs.init_new(server_args)
+        self.port_args = PortArgs.init_new(server_args)
         logger.info(f"{server_args=}")
 
         # Launch subprocesses
-        tokenizer_manager, scheduler_info = _launch_subprocesses(
+        tokenizer_manager, template_manager, scheduler_info = _launch_subprocesses(
             server_args=server_args,
-            port_args=port_args,
+            port_args=self.port_args,
         )
         self.server_args = server_args
         self.tokenizer_manager = tokenizer_manager
+        self.template_manager = template_manager
         self.scheduler_info = scheduler_info
 
         context = zmq.Context(2)
         self.send_to_rpc = get_zmq_socket(
-            context, zmq.DEALER, port_args.rpc_ipc_name, True
+            context, zmq.DEALER, self.port_args.rpc_ipc_name, True
         )
 
     def generate(
@@ -175,7 +172,7 @@ class Engine(EngineBase):
         """
         if self.server_args.enable_dp_attention:
             if data_parallel_rank is None:
-                logger.info("data_parallel_rank not provided, using default dispatch")
+                logger.debug("data_parallel_rank not provided, using default dispatch")
             elif data_parallel_rank < 0:
                 raise ValueError("data_parallel_rank must be non-negative")
             elif data_parallel_rank >= self.server_args.dp_size:
@@ -245,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,
@@ -258,7 +256,7 @@ class Engine(EngineBase):
 
         if self.server_args.enable_dp_attention:
             if data_parallel_rank is None:
-                logger.info("data_parallel_rank not provided, using default dispatch")
+                logger.debug("data_parallel_rank not provided, using default dispatch")
             elif data_parallel_rank < 0:
                 raise ValueError("data_parallel_rank must be non-negative")
             elif data_parallel_rank >= self.server_args.dp_size:
@@ -277,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,
@@ -479,17 +478,15 @@ class Engine(EngineBase):
             self.tokenizer_manager.get_weights_by_name(obj, None)
         )
 
-    def release_memory_occupation(self):
-        """Release GPU occupation temporarily."""
-        obj = ReleaseMemoryOccupationReqInput()
+    def release_memory_occupation(self, tags: Optional[List[str]] = None):
+        obj = ReleaseMemoryOccupationReqInput(tags=tags)
         loop = asyncio.get_event_loop()
         return loop.run_until_complete(
             self.tokenizer_manager.release_memory_occupation(obj, None)
         )
 
-    def resume_memory_occupation(self):
-        """Resume GPU occupation."""
-        obj = ResumeMemoryOccupationReqInput()
+    def resume_memory_occupation(self, tags: Optional[List[str]] = None):
+        obj = ResumeMemoryOccupationReqInput(tags=tags)
         loop = asyncio.get_event_loop()
         return loop.run_until_complete(
             self.tokenizer_manager.resume_memory_occupation(obj, None)
@@ -649,7 +646,7 @@ def _set_envs_and_config(server_args: ServerArgs):
 
 def _launch_subprocesses(
     server_args: ServerArgs, port_args: Optional[PortArgs] = None
-) -> Tuple[TokenizerManager, Dict]:
+) -> Tuple[TokenizerManager, TemplateManager, Dict]:
     """
     Launch the TokenizerManager in the main process, the Scheduler in a subprocess, and the DetokenizerManager in another subprocess.
     """
@@ -670,11 +667,9 @@ def _launch_subprocesses(
 
     scheduler_procs = []
     if server_args.dp_size == 1:
-        # Launch tensor parallel scheduler processes
         memory_saver_adapter = TorchMemorySaverAdapter.create(
             enable=server_args.enable_memory_saver
         )
-
         scheduler_pipe_readers = []
 
         nnodes_per_tp_group = max(server_args.nnodes // server_args.pp_size, 1)
@@ -710,6 +705,7 @@ def _launch_subprocesses(
                         writer,
                     ),
                 )
+
                 with memory_saver_adapter.configure_subprocess():
                     proc.start()
                 scheduler_procs.append(proc)
@@ -735,7 +731,7 @@ def _launch_subprocesses(
 
         if os.getenv("SGLANG_BLOCK_NONZERO_RANK_CHILDREN") == "0":
             # When using `Engine` as a Python API, we don't want to block here.
-            return None, None
+            return None, None, None
 
         launch_dummy_health_check_server(server_args.host, server_args.port)
 
@@ -744,7 +740,7 @@ def _launch_subprocesses(
             logger.error(
                 f"Scheduler or DataParallelController {proc.pid} terminated with {proc.exitcode}"
             )
-        return None, None
+        return None, None, None
 
     # Launch detokenizer process
     detoken_proc = mp.Process(
@@ -758,15 +754,15 @@ def _launch_subprocesses(
 
     # Launch tokenizer process
     tokenizer_manager = TokenizerManager(server_args, port_args)
-    if server_args.chat_template:
-        load_chat_template_for_openai_api(
-            tokenizer_manager, server_args.chat_template, server_args.model_path
-        )
-    else:
-        guess_chat_template_name_from_model_path(server_args.model_path)
 
-    if server_args.completion_template:
-        load_completion_template_for_openai_api(server_args.completion_template)
+    # Initialize templates
+    template_manager = TemplateManager()
+    template_manager.initialize_templates(
+        tokenizer_manager=tokenizer_manager,
+        model_path=server_args.model_path,
+        chat_template=server_args.chat_template,
+        completion_template=server_args.completion_template,
+    )
 
     # Wait for the model to finish loading
     scheduler_infos = []
@@ -790,4 +786,4 @@ def _launch_subprocesses(
     # Assume all schedulers have the same scheduler_info
     scheduler_info = scheduler_infos[0]
     tokenizer_manager.max_req_input_len = scheduler_info["max_req_input_len"]
-    return tokenizer_manager, scheduler_info
+    return tokenizer_manager, template_manager, scheduler_info