sglang 0.4.1.post5__py3-none-any.whl → 0.4.1.post7__py3-none-any.whl

sglang/srt/managers/session_controller.py

@@ -131,6 +131,7 @@ class Session:
             sampling_params=req.sampling_params,
             lora_path=req.lora_path,
             session_id=self.session_id,
+            custom_logit_processor=req.custom_logit_processor,
         )
         if last_req is not None:
             new_req.image_inputs = last_req.image_inputs

sglang/srt/managers/tokenizer_manager.py

@@ -18,10 +18,14 @@ import copy
 import dataclasses
 import logging
 import os
+import pickle
 import signal
 import sys
+import threading
 import time
 import uuid
+from datetime import datetime
+from http import HTTPStatus
 from typing import Any, Awaitable, Dict, Generic, List, Optional, Tuple, TypeVar, Union
 
 import fastapi
@@ -43,6 +47,7 @@ from sglang.srt.managers.io_struct import (
     BatchStrOut,
     BatchTokenIDOut,
     CloseSessionReqInput,
+    ConfigureLoggingReq,
     EmbeddingReqInput,
     FlushCacheReq,
     GenerateReqInput,
@@ -53,6 +58,10 @@ from sglang.srt.managers.io_struct import (
     OpenSessionReqInput,
     OpenSessionReqOutput,
     ProfileReq,
+    ReleaseMemoryOccupationReqInput,
+    ReleaseMemoryOccupationReqOutput,
+    ResumeMemoryOccupationReqInput,
+    ResumeMemoryOccupationReqOutput,
     SessionParams,
     TokenizedEmbeddingReqInput,
     TokenizedGenerateReqInput,
@@ -71,6 +80,7 @@ from sglang.srt.utils import (
     get_zmq_socket,
     kill_process_tree,
 )
+from sglang.utils import TypeBasedDispatcher, get_exception_traceback
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
@@ -103,16 +113,19 @@ class TokenizerManager:
         port_args: PortArgs,
     ):
         # Parse args
+
         self.server_args = server_args
         self.enable_metrics = server_args.enable_metrics
+        self.log_requests = server_args.log_requests
+        self.log_requests_level = 0
 
         # Init inter-process communication
         context = zmq.asyncio.Context(2)
         self.recv_from_detokenizer = get_zmq_socket(
-            context, zmq.PULL, port_args.tokenizer_ipc_name
+            context, zmq.PULL, port_args.tokenizer_ipc_name, True
         )
         self.send_to_scheduler = get_zmq_socket(
-            context, zmq.PUSH, port_args.scheduler_input_ipc_name
+            context, zmq.PUSH, port_args.scheduler_input_ipc_name, True
         )
 
         # Read model args
@@ -145,6 +158,7 @@ class TokenizerManager:
                     server_args.tokenizer_path,
                     tokenizer_mode=server_args.tokenizer_mode,
                     trust_remote_code=server_args.trust_remote_code,
+                    revision=server_args.revision,
                 )
                 self.tokenizer = self.processor.tokenizer
                 os.environ["TOKENIZERS_PARALLELISM"] = "false"
@@ -158,11 +172,15 @@ class TokenizerManager:
                     server_args.tokenizer_path,
                     tokenizer_mode=server_args.tokenizer_mode,
                     trust_remote_code=server_args.trust_remote_code,
+                    revision=server_args.revision,
                 )
 
         # Store states
-        self.to_create_loop = True
+        self.no_create_loop = False
         self.rid_to_state: Dict[str, ReqState] = {}
+        self.dump_requests_folder = ""  # By default do not dump
+        self.dump_requests_threshold = 1000
+        self.dump_request_list: List[Tuple] = []
 
         # The event to notify the weight sync is finished.
         self.model_update_lock = RWLock()
@@ -188,6 +206,14 @@ class TokenizerManager:
         self.get_weights_by_name_communicator = _Communicator(
             self.send_to_scheduler, server_args.dp_size
         )
+        self.release_memory_occupation_communicator = _Communicator(
+            self.send_to_scheduler, server_args.dp_size
+        )
+        self.resume_memory_occupation_communicator = _Communicator(
+            self.send_to_scheduler, server_args.dp_size
+        )
+        # Set after scheduler is initialized
+        self.max_req_input_len = None
 
         # Metrics
         if self.enable_metrics:
@@ -198,6 +224,44 @@ class TokenizerManager:
                 },
             )
 
+        self._result_dispatcher = TypeBasedDispatcher(
+            [
+                (
+                    (BatchStrOut, BatchEmbeddingOut, BatchTokenIDOut),
+                    self._handle_batch_output,
+                ),
+                (OpenSessionReqOutput, self._handle_open_session_req_output),
+                (
+                    UpdateWeightFromDiskReqOutput,
+                    self._handle_update_weights_from_disk_req_output,
+                ),
+                (
+                    InitWeightsUpdateGroupReqOutput,
+                    self.init_weights_update_group_communicator.handle_recv,
+                ),
+                (
+                    UpdateWeightsFromDistributedReqOutput,
+                    self.update_weights_from_distributed_communicator.handle_recv,
+                ),
+                (
+                    UpdateWeightsFromTensorReqOutput,
+                    self.update_weights_from_tensor_communicator.handle_recv,
+                ),
+                (
+                    GetWeightsByNameReqOutput,
+                    self.get_weights_by_name_communicator.handle_recv,
+                ),
+                (
+                    ReleaseMemoryOccupationReqOutput,
+                    self.release_memory_occupation_communicator.handle_recv,
+                ),
+                (
+                    ResumeMemoryOccupationReqOutput,
+                    self.resume_memory_occupation_communicator.handle_recv,
+                ),
+            ]
+        )
+
     async def generate_request(
         self,
         obj: Union[GenerateReqInput, EmbeddingReqInput],
@@ -215,8 +279,11 @@ class TokenizerManager:
 
         obj.normalize_batch_and_arguments()
 
-        if self.server_args.log_requests:
-            logger.info(f"Receive: obj={dataclass_to_string_truncated(obj)}")
+        if self.log_requests:
+            max_length = 2048 if self.log_requests_level == 0 else 1 << 30
+            logger.info(
+                f"Receive: obj={dataclass_to_string_truncated(obj, max_length)}"
+            )
 
         async with self.model_update_lock.reader_lock:
             is_single = obj.is_single
@@ -248,15 +315,21 @@ class TokenizerManager:
                 )
             input_embeds = obj.input_embeds
             input_ids = obj.input_ids
-        elif obj.input_ids is None:
-            input_ids = self.tokenizer.encode(input_text)
-        else:
+        elif obj.input_ids is not None:
             input_ids = obj.input_ids
+        else:
+            if self.tokenizer is None:
+                raise ValueError(
+                    "The engine initialized with skip_tokenizer_init=True cannot "
+                    "accept text prompts. Please provide input_ids or re-initialize "
+                    "the engine with skip_tokenizer_init=False."
+                )
+            input_ids = self.tokenizer.encode(input_text)
 
         if self.is_generation:
             # TODO: also support getting embeddings for multimodal models
             image_inputs: Dict = await self.image_processor.process_images_async(
-                obj.image_data, input_text or input_ids, obj
+                obj.image_data, input_text or input_ids, obj, self.max_req_input_len
             )
             if image_inputs and "input_ids" in image_inputs:
                 input_ids = image_inputs["input_ids"]
@@ -267,12 +340,28 @@ class TokenizerManager:
                 SessionParams(**obj.session_params) if obj.session_params else None
             )
 
-        if obj.input_ids is not None and len(input_ids) >= self.context_len:
+        input_token_num = len(input_ids) if input_ids is not None else 0
+        if input_token_num >= self.context_len:
             raise ValueError(
-                f"The input ({len(input_ids)} tokens) is longer than the "
+                f"The input ({input_token_num} tokens) is longer than the "
                 f"model's context length ({self.context_len} tokens)."
             )
 
+        if (
+            obj.sampling_params.get("max_new_tokens") is not None
+            and obj.sampling_params.get("max_new_tokens") + input_token_num
+            >= self.context_len
+        ):
+            raise ValueError(
+                f"Requested token count exceeds the model's maximum context length "
+                f"of {self.context_len} tokens. You requested a total of "
+                f"{obj.sampling_params.get('max_new_tokens') + input_token_num} "
+                f"tokens: {input_token_num} tokens from the input messages and "
+                f"{obj.sampling_params.get('max_new_tokens')} tokens for the "
+                f"completion. Please reduce the number of tokens in the input "
+                f"messages or the completion to fit within the limit."
+            )
+
         # Parse sampling parameters
         sampling_params = SamplingParams(**obj.sampling_params)
         sampling_params.normalize(self.tokenizer)
@@ -293,6 +382,7 @@ class TokenizerManager:
                 lora_path=obj.lora_path,
                 input_embeds=input_embeds,
                 session_params=session_params,
+                custom_logit_processor=obj.custom_logit_processor,
             )
         elif isinstance(obj, EmbeddingReqInput):
             tokenized_obj = TokenizedEmbeddingReqInput(
@@ -336,10 +426,21 @@ class TokenizerManager:
 
             state.out_list = []
             if state.finished:
-                if self.server_args.log_requests:
-                    msg = f"Finish: obj={dataclass_to_string_truncated(obj)}, out={dataclass_to_string_truncated(out)}"
+                if self.log_requests:
+                    max_length = 2048 if self.log_requests_level == 0 else 1 << 30
+                    msg = f"Finish: obj={dataclass_to_string_truncated(obj, max_length)}, out={dataclass_to_string_truncated(out, max_length)}"
                     logger.info(msg)
                 del self.rid_to_state[obj.rid]
+
+                # Check if this was an abort/error created by scheduler
+                if isinstance(out["meta_info"].get("finish_reason"), dict):
+                    finish_reason = out["meta_info"]["finish_reason"]
+                    if (
+                        finish_reason.get("type") == "abort"
+                        and finish_reason.get("status_code") == HTTPStatus.BAD_REQUEST
+                    ):
+                        raise ValueError(finish_reason["message"])
+
                 yield out
                 break
 
@@ -548,6 +649,22 @@ class TokenizerManager:
         else:
             return all_parameters
 
+    async def release_memory_occupation(
+        self,
+        obj: ReleaseMemoryOccupationReqInput,
+        request: Optional[fastapi.Request] = None,
+    ):
+        self.auto_create_handle_loop()
+        await self.release_memory_occupation_communicator(obj)
+
+    async def resume_memory_occupation(
+        self,
+        obj: ResumeMemoryOccupationReqInput,
+        request: Optional[fastapi.Request] = None,
+    ):
+        self.auto_create_handle_loop()
+        await self.resume_memory_occupation_communicator(obj)
+
     async def open_session(
         self, obj: OpenSessionReqInput, request: Optional[fastapi.Request] = None
     ):
@@ -568,9 +685,19 @@ class TokenizerManager:
     async def close_session(
         self, obj: CloseSessionReqInput, request: Optional[fastapi.Request] = None
     ):
-        assert not self.to_create_loop, "close session should not be the first request"
         await self.send_to_scheduler.send_pyobj(obj)
 
+    def configure_logging(self, obj: ConfigureLoggingReq):
+        if obj.log_requests is not None:
+            self.log_requests = obj.log_requests
+        if obj.log_requests_level is not None:
+            self.log_requests_level = obj.log_requests_level
+        if obj.dump_requests_folder is not None:
+            self.dump_requests_folder = obj.dump_requests_folder
+        if obj.dump_requests_threshold is not None:
+            self.dump_requests_threshold = obj.dump_requests_threshold
+        logging.info(f"Config logging: {obj=}")
+
     def create_abort_task(self, obj: GenerateReqInput):
         # Abort the request if the client is disconnected.
         async def abort_request():
@@ -586,22 +713,35 @@ class TokenizerManager:
         return background_tasks
 
     def auto_create_handle_loop(self):
-        if not self.to_create_loop:
+        if self.no_create_loop:
             return
 
-        self.to_create_loop = False
+        self.no_create_loop = True
         loop = asyncio.get_event_loop()
-        self.asyncio_tasks.add(loop.create_task(self.handle_loop()))
+        self.asyncio_tasks.add(
+            loop.create_task(print_exception_wrapper(self.handle_loop))
+        )
 
-        signal_handler = SignalHandler(self)
-        loop.add_signal_handler(signal.SIGTERM, signal_handler.signal_handler)
-        self.asyncio_tasks.add(loop.create_task(self.sigterm_watchdog()))
+        # We cannot add signal handler when the tokenizer manager is not in
+        # the main thread due to the CPython limitation.
+        if threading.current_thread() is threading.main_thread():
+            signal_handler = SignalHandler(self)
+            loop.add_signal_handler(signal.SIGTERM, signal_handler.signal_handler)
+        else:
+            logger.warning(
+                "Signal handler is not added because the tokenizer manager is "
+                "not in the main thread. This disables graceful shutdown of the "
+                "tokenizer manager when SIGTERM is received."
+            )
+        self.asyncio_tasks.add(
+            loop.create_task(print_exception_wrapper(self.sigterm_watchdog))
+        )
 
     async def sigterm_watchdog(self):
         while not self.gracefully_exit:
             await asyncio.sleep(5)
 
-        # drain requests
+        # Drain requests
         while True:
             remain_num_req = len(self.rid_to_state)
             logger.info(
@@ -619,143 +759,64 @@ class TokenizerManager:
         """The event loop that handles requests"""
 
         while True:
-            recv_obj: Union[
-                BatchStrOut,
-                BatchEmbeddingOut,
-                BatchTokenIDOut,
-                UpdateWeightFromDiskReqOutput,
-                UpdateWeightsFromDistributedReqOutput,
-                GetWeightsByNameReqOutput,
-                InitWeightsUpdateGroupReqOutput,
-            ] = await self.recv_from_detokenizer.recv_pyobj()
-
-            if isinstance(recv_obj, (BatchStrOut, BatchEmbeddingOut, BatchTokenIDOut)):
-                for i, rid in enumerate(recv_obj.rids):
-                    state = self.rid_to_state.get(rid, None)
-                    if state is None:
-                        continue
-
-                    meta_info = {
-                        "id": rid,
-                        "finish_reason": recv_obj.finished_reasons[i],
-                        "prompt_tokens": recv_obj.prompt_tokens[i],
-                    }
+            recv_obj = await self.recv_from_detokenizer.recv_pyobj()
+            self._result_dispatcher(recv_obj)
 
-                    if getattr(state.obj, "return_logprob", False):
-                        self.convert_logprob_style(
-                            meta_info,
-                            state.obj.top_logprobs_num,
-                            state.obj.return_text_in_logprobs,
-                            recv_obj,
-                            i,
-                        )
-
-                    if not isinstance(recv_obj, BatchEmbeddingOut):
-                        meta_info.update(
-                            {
-                                "completion_tokens": recv_obj.completion_tokens[i],
-                                "cached_tokens": recv_obj.cached_tokens[i],
-                            }
-                        )
-
-                    if isinstance(recv_obj, BatchStrOut):
-                        out_dict = {
-                            "text": recv_obj.output_strs[i],
-                            "meta_info": meta_info,
-                        }
-                        if self.server_args.return_token_ids:
-                            out_dict.update(
-                                {
-                                    "input_ids": recv_obj.origin_input_ids[i],
-                                    "output_ids": recv_obj.output_ids[i],
-                                }
-                            )
-                    elif isinstance(recv_obj, BatchTokenIDOut):
-                        out_dict = {
-                            "token_ids": recv_obj.output_ids[i],
-                            "meta_info": meta_info,
-                        }
-                    else:
-                        assert isinstance(recv_obj, BatchEmbeddingOut)
-                        out_dict = {
-                            "embedding": recv_obj.embeddings[i],
-                            "meta_info": meta_info,
-                        }
-                    state.out_list.append(out_dict)
-                    state.finished = recv_obj.finished_reasons[i] is not None
-                    state.event.set()
-
-                    if self.enable_metrics:
-                        completion_tokens = (
-                            recv_obj.completion_tokens[i]
-                            if getattr(recv_obj, "completion_tokens", None)
-                            else 0
-                        )
-
-                        if state.first_token_time is None:
-                            state.first_token_time = time.time()
-                            self.metrics_collector.observe_time_to_first_token(
-                                state.first_token_time - state.created_time
-                            )
-                        else:
-                            if completion_tokens >= 2:
-                                # Compute time_per_output_token for the streaming case
-                                self.metrics_collector.observe_time_per_output_token(
-                                    (time.time() - state.first_token_time)
-                                    / (completion_tokens - 1)
-                                )
-
-                        if state.finished:
-                            self.metrics_collector.inc_prompt_tokens(
-                                recv_obj.prompt_tokens[i]
-                            )
-                            self.metrics_collector.inc_generation_tokens(
-                                completion_tokens
-                            )
-                            self.metrics_collector.observe_e2e_request_latency(
-                                time.time() - state.created_time
-                            )
-                            # Compute time_per_output_token for the non-streaming case
-                            if (
-                                hasattr(state.obj, "stream")
-                                and not state.obj.stream
-                                and completion_tokens >= 1
-                            ):
-                                self.metrics_collector.observe_time_per_output_token(
-                                    (time.time() - state.created_time)
-                                    / completion_tokens
-                                )
-            elif isinstance(recv_obj, OpenSessionReqOutput):
-                self.session_futures[recv_obj.session_id].set_result(
-                    recv_obj.session_id if recv_obj.success else None
+    def _handle_batch_output(
+        self, recv_obj: Union[BatchStrOut, BatchEmbeddingOut, BatchTokenIDOut]
+    ):
+        for i, rid in enumerate(recv_obj.rids):
+            state = self.rid_to_state.get(rid, None)
+            if state is None:
+                continue
+
+            meta_info = {
+                "id": rid,
+                "finish_reason": recv_obj.finished_reasons[i],
+                "prompt_tokens": recv_obj.prompt_tokens[i],
+            }
+
+            if getattr(state.obj, "return_logprob", False):
+                self.convert_logprob_style(
+                    meta_info,
+                    state.obj.top_logprobs_num,
+                    state.obj.return_text_in_logprobs,
+                    recv_obj,
+                    i,
+                )
+
+            if not isinstance(recv_obj, BatchEmbeddingOut):
+                meta_info.update(
+                    {
+                        "completion_tokens": recv_obj.completion_tokens[i],
+                        "cached_tokens": recv_obj.cached_tokens[i],
+                    }
                 )
-            elif isinstance(recv_obj, UpdateWeightFromDiskReqOutput):
-                if self.server_args.dp_size == 1:
-                    self.model_update_result.set_result(recv_obj)
-                else:  # self.server_args.dp_size > 1
-                    self.model_update_tmp.append(recv_obj)
-                    # set future if the all results are recevied
-                    if len(self.model_update_tmp) == self.server_args.dp_size:
-                        self.model_update_result.set_result(self.model_update_tmp)
-            elif isinstance(recv_obj, InitWeightsUpdateGroupReqOutput):
-                assert (
-                    self.server_args.dp_size == 1
-                ), "dp_size must be 1 for init parameter update group"
-                self.init_weights_update_group_communicator.handle_recv(recv_obj)
-            elif isinstance(recv_obj, UpdateWeightsFromDistributedReqOutput):
-                assert (
-                    self.server_args.dp_size == 1
-                ), "dp_size must be 1 for update weights from distributed"
-                self.update_weights_from_distributed_communicator.handle_recv(recv_obj)
-            elif isinstance(recv_obj, UpdateWeightsFromTensorReqOutput):
-                assert (
-                    self.server_args.dp_size == 1
-                ), "dp_size must be 1 for update weights from distributed"
-                self.update_weights_from_tensor_communicator.handle_recv(recv_obj)
-            elif isinstance(recv_obj, GetWeightsByNameReqOutput):
-                self.get_weights_by_name_communicator.handle_recv(recv_obj)
+
+            if isinstance(recv_obj, BatchStrOut):
+                out_dict = {
+                    "text": recv_obj.output_strs[i],
+                    "meta_info": meta_info,
+                }
+            elif isinstance(recv_obj, BatchTokenIDOut):
+                out_dict = {
+                    "token_ids": recv_obj.output_ids[i],
+                    "meta_info": meta_info,
+                }
             else:
-                raise ValueError(f"Invalid object: {recv_obj=}")
+                assert isinstance(recv_obj, BatchEmbeddingOut)
+                out_dict = {
+                    "embedding": recv_obj.embeddings[i],
+                    "meta_info": meta_info,
+                }
+            state.out_list.append(out_dict)
+            state.finished = recv_obj.finished_reasons[i] is not None
+            state.event.set()
+
+            if self.enable_metrics and state.obj.log_metrics:
+                self.collect_metrics(state, recv_obj, i)
+            if self.dump_requests_folder and state.finished and state.obj.log_metrics:
+                self.dump_requests(state, out_dict)
 
     def convert_logprob_style(
         self,
@@ -775,9 +836,6 @@ class TokenizerManager:
             recv_obj.output_token_logprobs_idx[recv_obj_index],
             return_text_in_logprobs,
         )
-        meta_info["normalized_prompt_logprob"] = recv_obj.normalized_prompt_logprob[
-            recv_obj_index
-        ]
 
         if top_logprobs_num > 0:
             meta_info["input_top_logprobs"] = self.detokenize_top_logprobs_tokens(
@@ -827,6 +885,93 @@ class TokenizerManager:
                 ret.append(None)
         return ret
 
+    def collect_metrics(self, state: ReqState, recv_obj: BatchStrOut, i: int):
+        completion_tokens = (
+            recv_obj.completion_tokens[i]
+            if getattr(recv_obj, "completion_tokens", None)
+            else 0
+        )
+
+        if state.first_token_time is None:
+            state.first_token_time = time.time()
+            self.metrics_collector.observe_time_to_first_token(
+                state.first_token_time - state.created_time
+            )
+        else:
+            if completion_tokens >= 2:
+                # Compute time_per_output_token for the streaming case
+                self.metrics_collector.observe_time_per_output_token(
+                    (time.time() - state.first_token_time) / (completion_tokens - 1)
+                )
+
+        if state.finished:
+            self.metrics_collector.observe_one_finished_request(
+                recv_obj.prompt_tokens[i], completion_tokens
+            )
+            self.metrics_collector.observe_e2e_request_latency(
+                time.time() - state.created_time
+            )
+            # Compute time_per_output_token for the non-streaming case
+            if (
+                hasattr(state.obj, "stream")
+                and not state.obj.stream
+                and completion_tokens >= 1
+            ):
+                self.metrics_collector.observe_time_per_output_token(
+                    (time.time() - state.created_time) / completion_tokens
+                )
+
+    def dump_requests(self, state: ReqState, out_dict: dict):
+        self.dump_request_list.append(
+            (state.obj, out_dict, state.created_time, time.time())
+        )
+
+        if len(self.dump_request_list) >= self.dump_requests_threshold:
+            filename = os.path.join(
+                self.dump_requests_folder,
+                datetime.now().strftime("%Y-%m-%d_%H-%M-%S") + ".pkl",
+            )
+            logger.info(f"Dump {len(self.dump_request_list)} requests to {filename}")
+
+            to_dump = self.dump_request_list
+            self.dump_request_list = []
+
+            def background_task():
+                os.makedirs(self.dump_requests_folder, exist_ok=True)
+                with open(filename, "wb") as f:
+                    pickle.dump(to_dump, f)
+
+            # Schedule the task to run in the background without awaiting it
+            asyncio.create_task(asyncio.to_thread(background_task))
+
+    def _handle_open_session_req_output(self, recv_obj):
+        self.session_futures[recv_obj.session_id].set_result(
+            recv_obj.session_id if recv_obj.success else None
+        )
+
+    def _handle_update_weights_from_disk_req_output(self, recv_obj):
+        if self.server_args.dp_size == 1:
+            self.model_update_result.set_result(recv_obj)
+        else:  # self.server_args.dp_size > 1
+            self.model_update_tmp.append(recv_obj)
+            # set future if the all results are recevied
+            if len(self.model_update_tmp) == self.server_args.dp_size:
+                self.model_update_result.set_result(self.model_update_tmp)
+
+
+async def print_exception_wrapper(func):
+    """
+    Sometimes an asyncio function does not print exception.
+    We do another wrapper to handle the exception.
+    """
+    try:
+        await func()
+    except Exception:
+        traceback = get_exception_traceback()
+        logger.error(f"TokenizerManager hit an exception: {traceback}")
+        kill_process_tree(os.getpid(), include_parent=True)
+        sys.exit(1)
+
 
 class SignalHandler:
     def __init__(self, tokenizer_manager):

sglang/srt/managers/tp_worker.py

@@ -83,6 +83,7 @@ class TpModelWorker:
                     server_args.tokenizer_path,
                     tokenizer_mode=server_args.tokenizer_mode,
                     trust_remote_code=server_args.trust_remote_code,
+                    revision=server_args.revision,
                 )
                 self.tokenizer = self.processor.tokenizer
             else:
@@ -90,6 +91,7 @@ class TpModelWorker:
                     server_args.tokenizer_path,
                     tokenizer_mode=server_args.tokenizer_mode,
                     trust_remote_code=server_args.trust_remote_code,
+                    revision=server_args.revision,
                 )
         self.device = self.model_runner.device
 
@@ -101,6 +103,7 @@ class TpModelWorker:
                 self.max_total_num_tokens // 2
                 if server_args.max_running_requests is None
                 else server_args.max_running_requests
+                // (server_args.dp_size if server_args.enable_dp_attention else 1)
             ),
             self.model_runner.req_to_token_pool.size,
         )
@@ -142,16 +145,15 @@ class TpModelWorker:
     def get_tp_cpu_group(self):
         return self.model_runner.tp_group.cpu_group
 
+    def get_attention_tp_cpu_group(self):
+        return self.model_runner.attention_tp_group.cpu_group
+
     def get_memory_pool(self):
         return (
             self.model_runner.req_to_token_pool,
             self.model_runner.token_to_kv_pool,
         )
 
-    def forward_batch_idle(self, model_worker_batch: ModelWorkerBatch):
-        forward_batch = ForwardBatch.init_new(model_worker_batch, self.model_runner)
-        self.model_runner.forward(forward_batch)
-
     def forward_batch_generation(
         self,
         model_worker_batch: ModelWorkerBatch,