sglang 0.5.2rc2__py3-none-any.whl → 0.5.3__py3-none-any.whl

sglang/srt/managers/io_struct.py

@@ -18,6 +18,7 @@ processes (TokenizerManager, DetokenizerManager, Scheduler).
 
 import copy
 import uuid
+from abc import ABC
 from dataclasses import dataclass, field
 from enum import Enum
 from typing import TYPE_CHECKING, Any, Dict, List, Optional, Union
@@ -35,10 +36,33 @@ else:
     Image = Any
 
 
+# Parameters for a session
+@dataclass
+class BaseReq(ABC):
+    rid: Optional[Union[str, List[str]]] = field(default=None, kw_only=True)
+
+    def regenerate_rid(self):
+        """Generate a new request ID and return it."""
+        if isinstance(self.rid, list):
+            self.rid = [uuid.uuid4().hex for _ in range(len(self.rid))]
+        else:
+            self.rid = uuid.uuid4().hex
+        return self.rid
+
+
+@dataclass
+class BaseBatchReq(ABC):
+    rids: Optional[List[str]] = field(default=None, kw_only=True)
+
+    def regenerate_rids(self):
+        """Generate new request IDs and return them."""
+        self.rids = [uuid.uuid4().hex for _ in range(len(self.rids))]
+        return self.rids
+
+
 @dataclass
 class SessionParams:
     id: Optional[str] = None
-    rid: Optional[str] = None
     offset: Optional[int] = None
     replace: Optional[bool] = None
     drop_previous_output: Optional[bool] = None
@@ -62,7 +86,7 @@ MultimodalDataInputFormat = Union[
 
 
 @dataclass
-class GenerateReqInput:
+class GenerateReqInput(BaseReq):
     # The input prompt. It can be a single prompt or a batch of prompts.
     text: Optional[Union[List[str], str]] = None
     # The token ids for text; one can specify either text or input_ids
@@ -82,8 +106,6 @@ class GenerateReqInput:
     audio_data: Optional[MultimodalDataInputFormat] = None
     # The sampling_params. See descriptions below.
     sampling_params: Optional[Union[List[Dict], Dict]] = None
-    # The request id.
-    rid: Optional[Union[List[str], str]] = None
     # Whether to return logprobs.
     return_logprob: Optional[Union[List[bool], bool]] = None
     # If return logprobs, the start location in the prompt for returning logprobs.
@@ -121,6 +143,7 @@ class GenerateReqInput:
     bootstrap_host: Optional[Union[List[str], str]] = None
     bootstrap_port: Optional[Union[List[Optional[int]], int]] = None
     bootstrap_room: Optional[Union[List[int], int]] = None
+    bootstrap_pair_key: Optional[Union[List[str], str]] = None
 
     # For data parallel rank routing
     data_parallel_rank: Optional[int] = None
@@ -128,6 +151,24 @@ class GenerateReqInput:
     # For background responses (OpenAI responses API)
     background: bool = False
 
+    # Conversation id used for tracking requests
+    conversation_id: Optional[str] = None
+
+    # Priority for the request
+    priority: Optional[int] = None
+
+    # Extra key for classifying the request (e.g. cache_salt)
+    extra_key: Optional[Union[List[str], str]] = None
+
+    # Whether to disallow logging for this request (e.g. due to ZDR)
+    no_logs: bool = False
+
+    # For custom metric labels
+    custom_labels: Optional[Dict[str, str]] = None
+
+    # (Internal) Whether to return bytes for image generation
+    return_bytes: bool = False
+
     def contains_mm_input(self) -> bool:
         return (
             has_valid_data(self.image_data)
@@ -258,6 +299,7 @@ class GenerateReqInput:
         self._normalize_sampling_params(num)
         self._normalize_logprob_params(num)
         self._normalize_custom_logit_processor(num)
+        self._normalize_bootstrap_params(num)
 
     def _expand_inputs(self, num):
         """Expand the main inputs (text, input_ids, input_embeds) for parallel sampling."""
@@ -297,6 +339,11 @@ class GenerateReqInput:
             self.image_data = [[self.image_data]] * num
             self.modalities = ["image"] * num
         elif isinstance(self.image_data, list):
+            # Handle empty list case - treat as no images
+            if len(self.image_data) == 0:
+                self.image_data = [None] * num
+                return
+
             if len(self.image_data) != self.batch_size:
                 raise ValueError(
                     "The length of image_data should be equal to the batch size."
@@ -421,6 +468,40 @@ class GenerateReqInput:
                 "Cannot use list custom_logit_processor with parallel_sample_num > 1"
             )
 
+    def _normalize_bootstrap_params(self, num):
+        """Normalize bootstrap parameters for batch processing."""
+        # Normalize bootstrap_host
+        if self.bootstrap_host is None:
+            self.bootstrap_host = [None] * num
+        elif not isinstance(self.bootstrap_host, list):
+            self.bootstrap_host = [self.bootstrap_host] * num
+        elif isinstance(self.bootstrap_host, list):
+            self.bootstrap_host = self.bootstrap_host * self.parallel_sample_num
+
+        # Normalize bootstrap_port
+        if self.bootstrap_port is None:
+            self.bootstrap_port = [None] * num
+        elif not isinstance(self.bootstrap_port, list):
+            self.bootstrap_port = [self.bootstrap_port] * num
+        elif isinstance(self.bootstrap_port, list):
+            self.bootstrap_port = self.bootstrap_port * self.parallel_sample_num
+
+        # Normalize bootstrap_room
+        if self.bootstrap_room is None:
+            self.bootstrap_room = [None] * num
+        elif not isinstance(self.bootstrap_room, list):
+            self.bootstrap_room = [self.bootstrap_room + i for i in range(num)]
+        elif isinstance(self.bootstrap_room, list):
+            self.bootstrap_room = self.bootstrap_room * self.parallel_sample_num
+
+        # Normalize bootstrap_pair_key
+        if self.bootstrap_pair_key is None:
+            self.bootstrap_pair_key = [None] * num
+        elif not isinstance(self.bootstrap_pair_key, list):
+            self.bootstrap_pair_key = [self.bootstrap_pair_key] * num
+        elif isinstance(self.bootstrap_pair_key, list):
+            self.bootstrap_pair_key = self.bootstrap_pair_key * self.parallel_sample_num
+
     def _validate_session_params(self):
         """Validate that session parameters are properly formatted."""
         if self.session_params is not None:
@@ -429,11 +510,6 @@ class GenerateReqInput:
             ):
                 raise ValueError("Session params must be a dict or a list of dicts.")
 
-    def regenerate_rid(self):
-        """Generate a new request ID and return it."""
-        self.rid = uuid.uuid4().hex
-        return self.rid
-
     def __getitem__(self, i):
         return GenerateReqInput(
             text=self.text[i] if self.text is not None else None,
@@ -453,7 +529,13 @@ class GenerateReqInput:
             return_text_in_logprobs=self.return_text_in_logprobs,
             stream=self.stream,
             log_metrics=self.log_metrics,
+            return_hidden_states=(
+                self.return_hidden_states[i]
+                if isinstance(self.return_hidden_states, list)
+                else self.return_hidden_states
+            ),
             modalities=self.modalities[i] if self.modalities else None,
+            session_params=self.session_params,
             lora_path=self.lora_path[i] if self.lora_path is not None else None,
             lora_id=self.lora_id[i] if self.lora_id is not None else None,
             custom_logit_processor=(
@@ -461,11 +543,6 @@ class GenerateReqInput:
                 if self.custom_logit_processor is not None
                 else None
             ),
-            return_hidden_states=(
-                self.return_hidden_states[i]
-                if isinstance(self.return_hidden_states, list)
-                else self.return_hidden_states
-            ),
             # if `__getitem__` is called, the bootstrap_host, bootstrap_port, bootstrap_room must be a list
             bootstrap_host=(
                 self.bootstrap_host[i] if self.bootstrap_host is not None else None
@@ -476,16 +553,25 @@ class GenerateReqInput:
             bootstrap_room=(
                 self.bootstrap_room[i] if self.bootstrap_room is not None else None
             ),
+            bootstrap_pair_key=(
+                self.bootstrap_pair_key[i]
+                if self.bootstrap_pair_key is not None
+                else None
+            ),
             data_parallel_rank=(
                 self.data_parallel_rank if self.data_parallel_rank is not None else None
             ),
+            conversation_id=self.conversation_id,
+            priority=self.priority,
+            extra_key=self.extra_key,
+            no_logs=self.no_logs,
+            custom_labels=self.custom_labels,
+            return_bytes=self.return_bytes,
         )
 
 
 @dataclass
-class TokenizedGenerateReqInput:
-    # The request id
-    rid: str
+class TokenizedGenerateReqInput(BaseReq):
     # The input text
     input_text: str
     # The input token ids
@@ -505,36 +591,50 @@ class TokenizedGenerateReqInput:
     # Whether to stream output
     stream: bool
 
-    # LoRA related
-    lora_id: Optional[str] = None  # None means just use the base model
+    # Whether to return hidden states
+    return_hidden_states: bool = False
+
     # The input embeds
     input_embeds: Optional[Union[List[List[List[float]]], List[List[float]]]] = None
 
     # Session info for continual prompting
     session_params: Optional[SessionParams] = None
 
+    # LoRA related
+    lora_id: Optional[str] = None  # None means just use the base model
+
     # Custom logit processor for advanced sampling control. Must be a serialized instance
     # of `CustomLogitProcessor` in python/sglang/srt/sampling/custom_logit_processor.py
     # Use the processor's `to_str()` method to generate the serialized string.
     custom_logit_processor: Optional[str] = None
 
-    # Whether to return hidden states
-    return_hidden_states: bool = False
-
     # For disaggregated inference
     bootstrap_host: Optional[str] = None
     bootstrap_port: Optional[int] = None
     bootstrap_room: Optional[int] = None
+    bootstrap_pair_key: Optional[str] = None
 
     # For data parallel rank routing
     data_parallel_rank: Optional[int] = None
 
-    # For dp balance
-    dp_balance_id: int = -1
+    # Priority for the request
+    priority: Optional[int] = None
+
+    # Extra key for classifying the request (e.g. cache_salt)
+    extra_key: Optional[str] = None
+
+    # Whether to disallow logging for this request (e.g. due to ZDR)
+    no_logs: bool = False
+
+    # tracing context
+    trace_context: Optional[Dict] = None
+
+    # (Internal) Whether to return bytes for image generation
+    return_bytes: bool = False
 
 
 @dataclass
-class BatchTokenizedGenerateReqInput:
+class BatchTokenizedGenerateReqInput(BaseBatchReq):
     # The batch of tokenized requests
     batch: List[TokenizedGenerateReqInput]
 
@@ -549,7 +649,7 @@ class BatchTokenizedGenerateReqInput:
 
 
 @dataclass
-class EmbeddingReqInput:
+class EmbeddingReqInput(BaseReq):
     # The input prompt. It can be a single prompt or a batch of prompts.
     text: Optional[Union[List[List[str]], List[str], str]] = None
     # The image input. It can be an image instance, file name, URL, or base64 encoded string.
@@ -565,8 +665,6 @@ class EmbeddingReqInput:
     audio_data: Optional[MultimodalDataInputFormat] = None
     # The token ids for text; one can either specify text or input_ids.
     input_ids: Optional[Union[List[List[int]], List[int]]] = None
-    # The request id.
-    rid: Optional[Union[List[str], str]] = None
     # Dummy sampling params for compatibility
     sampling_params: Optional[Union[List[Dict], Dict]] = None
     # Dummy input embeds for compatibility
@@ -577,10 +675,15 @@ class EmbeddingReqInput:
     modalities: Optional[List[str]] = None
     # For cross-encoder requests
     is_cross_encoder_request: bool = False
+    # Priority for the request
+    priority: Optional[int] = None
 
     # For background responses (OpenAI responses API)
     background: bool = False
 
+    # tracing context
+    trace_context: Optional[Dict] = None
+
     def normalize_batch_and_arguments(self):
         # at least one of text, input_ids, or image should be provided
         if self.text is None and self.input_ids is None and self.image_data is None:
@@ -632,10 +735,6 @@ class EmbeddingReqInput:
             for i in range(self.batch_size):
                 self.sampling_params[i]["max_new_tokens"] = 0
 
-    def regenerate_rid(self):
-        self.rid = uuid.uuid4().hex
-        return self.rid
-
     def contains_mm_input(self) -> bool:
         return (
             has_valid_data(self.image_data)
@@ -664,9 +763,7 @@ class EmbeddingReqInput:
 
 
 @dataclass
-class TokenizedEmbeddingReqInput:
-    # The request id
-    rid: str
+class TokenizedEmbeddingReqInput(BaseReq):
     # The input text
     input_text: str
     # The input token ids
@@ -679,12 +776,12 @@ class TokenizedEmbeddingReqInput:
     sampling_params: SamplingParams
     # For data parallel rank routing
     data_parallel_rank: Optional[int] = None
-    # For dp balance
-    dp_balance_id: int = -1
+    # Priority for the request
+    priority: Optional[int] = None
 
 
 @dataclass
-class BatchTokenizedEmbeddingReqInput:
+class BatchTokenizedEmbeddingReqInput(BaseBatchReq):
     # The batch of tokenized embedding requests
     batch: List[TokenizedEmbeddingReqInput]
 
@@ -699,9 +796,7 @@ class BatchTokenizedEmbeddingReqInput:
 
 
 @dataclass
-class BatchTokenIDOut:
-    # The request id
-    rids: List[str]
+class BatchTokenIDOutput(BaseBatchReq):
     # The finish reason
     finished_reasons: List[BaseFinishReason]
     # For incremental decoding
@@ -738,11 +833,26 @@ class BatchTokenIDOut:
     # Hidden states
     output_hidden_states: List[List[float]]
 
+    # The information of placeholder tokens (e.g., image token)
+    # idx is the index of the token in the prompt after expansion.
+    # val is the length of padded tokens after expansion.
+    placeholder_tokens_idx: List[Optional[List[int]]]
+    placeholder_tokens_val: List[Optional[List[int]]]
+
 
 @dataclass
-class BatchMultimodalDecodeReq:
-    # The request id
-    rids: List[str]
+class BatchMultimodalDecodeReq(BaseBatchReq):
+    decoded_ids: List[int]
+    input_token_logprobs_val: List[float]
+    input_token_logprobs_idx: List[int]
+    output_token_logprobs_val: List[float]
+    output_token_logprobs_idx: List[int]
+    read_offsets: List[int]
+    skip_special_tokens: List[bool]
+    spaces_between_special_tokens: List[bool]
+    image_resolutions: List[List[int]]
+    resize_image_resolutions: List[List[int]]
+
     finished_reasons: List[BaseFinishReason]
 
     # Token counts
@@ -750,11 +860,15 @@ class BatchMultimodalDecodeReq:
     completion_tokens: List[int]
     cached_tokens: List[int]
 
+    # Placeholder token info
+    placeholder_tokens_idx: List[Optional[List[int]]]
+    placeholder_tokens_val: List[Optional[List[int]]]
+
+    return_bytes: bool = False
+
 
 @dataclass
-class BatchStrOut:
-    # The request id
-    rids: List[str]
+class BatchStrOutput(BaseBatchReq):
     # The finish reason
     finished_reasons: List[dict]
     # The output decoded strings
@@ -785,26 +899,37 @@ class BatchStrOut:
     # Hidden states
     output_hidden_states: List[List[float]]
 
+    placeholder_tokens_idx: List[Optional[List[int]]]
+    placeholder_tokens_val: List[Optional[List[int]]]
+
 
 @dataclass
-class BatchMultimodalOut:
-    # The request id
-    rids: List[str]
+class BatchMultimodalOutput(BaseBatchReq):
     # The finish reason
     finished_reasons: List[dict]
+    decoded_ids: List[List[int]]
     # The outputs
-    outputs: List[List[Dict]]
+    outputs: Union[List[str | bytes], List[List[Dict]]]
+
+    # probability values for input tokens and output tokens
+    input_token_logprobs_val: List[List[float]]
+    input_token_logprobs_idx: List[List[int]]
+    output_token_logprobs_val: List[List[float]]
+    output_token_logprobs_idx: List[List[int]]
 
     # Token counts
     prompt_tokens: List[int]
     completion_tokens: List[int]
     cached_tokens: List[int]
 
+    placeholder_tokens_idx: List[Optional[List[int]]]
+    placeholder_tokens_val: List[Optional[List[int]]]
+
+    return_bytes: List[bool]
+
 
 @dataclass
-class BatchEmbeddingOut:
-    # The request id
-    rids: List[str]
+class BatchEmbeddingOutput(BaseBatchReq):
     # The finish reason
     finished_reasons: List[BaseFinishReason]
     # The output embedding
@@ -812,30 +937,33 @@ class BatchEmbeddingOut:
     # Token counts
     prompt_tokens: List[int]
     cached_tokens: List[int]
+    # Placeholder token info
+    placeholder_tokens_idx: List[Optional[List[int]]]
+    placeholder_tokens_val: List[Optional[List[int]]]
 
 
 @dataclass
-class ClearHiCacheReqInput:
+class ClearHiCacheReqInput(BaseReq):
     pass
 
 
 @dataclass
-class ClearHiCacheReqOutput:
+class ClearHiCacheReqOutput(BaseReq):
     success: bool
 
 
 @dataclass
-class FlushCacheReqInput:
+class FlushCacheReqInput(BaseReq):
     pass
 
 
 @dataclass
-class FlushCacheReqOutput:
+class FlushCacheReqOutput(BaseReq):
     success: bool
 
 
 @dataclass
-class UpdateWeightFromDiskReqInput:
+class UpdateWeightFromDiskReqInput(BaseReq):
     # The model path with the new weights
     model_path: str
     # The format to load the weights
@@ -844,10 +972,16 @@ class UpdateWeightFromDiskReqInput:
     abort_all_requests: bool = False
     # Optional: Update weight version along with weights
     weight_version: Optional[str] = None
+    # Whether to update weights asynchronously
+    is_async: bool = False
+    # Whether to empty torch cache
+    torch_empty_cache: bool = False
+    # Whether to keep the scheduler paused after weight update
+    keep_pause: bool = False
 
 
 @dataclass
-class UpdateWeightFromDiskReqOutput:
+class UpdateWeightFromDiskReqOutput(BaseReq):
     success: bool
     message: str
     # Number of paused requests during weight sync.
@@ -855,7 +989,7 @@ class UpdateWeightFromDiskReqOutput:
 
 
 @dataclass
-class UpdateWeightsFromDistributedReqInput:
+class UpdateWeightsFromDistributedReqInput(BaseReq):
     names: List[str]
     dtypes: List[str]
     shapes: List[List[int]]
@@ -870,13 +1004,13 @@ class UpdateWeightsFromDistributedReqInput:
 
 
 @dataclass
-class UpdateWeightsFromDistributedReqOutput:
+class UpdateWeightsFromDistributedReqOutput(BaseReq):
     success: bool
     message: str
 
 
 @dataclass
-class UpdateWeightsFromTensorReqInput:
+class UpdateWeightsFromTensorReqInput(BaseReq):
     """Update model weights from tensor input.
 
     - Tensors are serialized for transmission
@@ -895,13 +1029,51 @@ class UpdateWeightsFromTensorReqInput:
 
 
 @dataclass
-class UpdateWeightsFromTensorReqOutput:
+class UpdateWeightsFromTensorReqOutput(BaseReq):
+    success: bool
+    message: str
+
+
+@dataclass
+class InitWeightsSendGroupForRemoteInstanceReqInput(BaseReq):
+    # The master address
+    master_address: str
+    # The ports for each rank's communication group
+    ports: str
+    # The rank in the communication group
+    group_rank: int
+    # The world size
+    world_size: int
+    # The group name
+    group_name: str = "weight_send_group"
+    # The backend
+    backend: str = "nccl"
+
+
+@dataclass
+class InitWeightsSendGroupForRemoteInstanceReqOutput(BaseReq):
+    success: bool
+    message: str
+
+
+@dataclass
+class SendWeightsToRemoteInstanceReqInput(BaseReq):
+    # The master address
+    master_address: str
+    # The ports for each rank's communication group
+    ports: str
+    # The group name
+    group_name: str = "weight_send_group"
+
+
+@dataclass
+class SendWeightsToRemoteInstanceReqOutput(BaseReq):
     success: bool
     message: str
 
 
 @dataclass
-class InitWeightsUpdateGroupReqInput:
+class InitWeightsUpdateGroupReqInput(BaseReq):
     # The master address
     master_address: str
     # The master port
@@ -917,13 +1089,24 @@ class InitWeightsUpdateGroupReqInput:
 
 
 @dataclass
-class InitWeightsUpdateGroupReqOutput:
+class InitWeightsUpdateGroupReqOutput(BaseReq):
+    success: bool
+    message: str
+
+
+@dataclass
+class DestroyWeightsUpdateGroupReqInput(BaseReq):
+    group_name: str = "weight_update_group"
+
+
+@dataclass
+class DestroyWeightsUpdateGroupReqOutput(BaseReq):
     success: bool
     message: str
 
 
 @dataclass
-class UpdateWeightVersionReqInput:
+class UpdateWeightVersionReqInput(BaseReq):
     # The new weight version
     new_version: str
     # Whether to abort all running requests before updating
@@ -931,88 +1114,87 @@ class UpdateWeightVersionReqInput:
 
 
 @dataclass
-class GetWeightsByNameReqInput:
+class GetWeightsByNameReqInput(BaseReq):
     name: str
     truncate_size: int = 100
 
 
 @dataclass
-class GetWeightsByNameReqOutput:
+class GetWeightsByNameReqOutput(BaseReq):
     parameter: list
 
 
 @dataclass
-class ReleaseMemoryOccupationReqInput:
+class ReleaseMemoryOccupationReqInput(BaseReq):
     # Optional tags to identify the memory region, which is primarily used for RL
     # Currently we only support `weights` and `kv_cache`
     tags: Optional[List[str]] = None
 
 
 @dataclass
-class ReleaseMemoryOccupationReqOutput:
+class ReleaseMemoryOccupationReqOutput(BaseReq):
     pass
 
 
 @dataclass
-class ResumeMemoryOccupationReqInput:
+class ResumeMemoryOccupationReqInput(BaseReq):
     # Optional tags to identify the memory region, which is primarily used for RL
     # Currently we only support `weights` and `kv_cache`
     tags: Optional[List[str]] = None
 
 
 @dataclass
-class ResumeMemoryOccupationReqOutput:
+class ResumeMemoryOccupationReqOutput(BaseReq):
     pass
 
 
 @dataclass
-class SlowDownReqInput:
+class SlowDownReqInput(BaseReq):
     forward_sleep_time: Optional[float]
 
 
 @dataclass
-class SlowDownReqOutput:
+class SlowDownReqOutput(BaseReq):
     pass
 
 
 @dataclass
-class AbortReq:
-    # The request id
-    rid: str = ""
+class AbortReq(BaseReq):
     # Whether to abort all requests
     abort_all: bool = False
     # The finished reason data
     finished_reason: Optional[Dict[str, Any]] = None
-    # used in MultiTokenzierManager mode
-    rids: Optional[Union[List[str], str]] = None
+    abort_reason: Optional[str] = None
 
     def __post_init__(self):
-        self.rids = self.rid
+        # FIXME: This is a hack to keep the same with the old code
+        if self.rid is None:
+            self.rid = ""
 
 
 @dataclass
-class GetInternalStateReq:
+class GetInternalStateReq(BaseReq):
     pass
 
 
 @dataclass
-class GetInternalStateReqOutput:
+class GetInternalStateReqOutput(BaseReq):
     internal_state: Dict[Any, Any]
 
 
 @dataclass
-class SetInternalStateReq:
+class SetInternalStateReq(BaseReq):
     server_args: Dict[str, Any]
 
 
 @dataclass
-class SetInternalStateReqOutput:
+class SetInternalStateReqOutput(BaseReq):
     updated: bool
     server_args: Dict[str, Any]
 
 
 @dataclass
-class ProfileReqInput:
+class ProfileReqInput(BaseReq):
     # The output directory
     output_dir: Optional[str] = None
     # If set, it profile as many as this number of steps.
@@ -1032,7 +1214,7 @@ class ProfileReqType(Enum):
 
 
 @dataclass
-class ProfileReq:
+class ProfileReq(BaseReq):
     type: ProfileReqType
     output_dir: Optional[str] = None
     start_step: Optional[int] = None
@@ -1045,54 +1227,59 @@ class ProfileReq:
 
 
 @dataclass
-class ProfileReqOutput:
+class ProfileReqOutput(BaseReq):
     success: bool
     message: str
 
 
 @dataclass
-class FreezeGCReq:
+class FreezeGCReq(BaseReq):
     pass
 
 
 @dataclass
-class ConfigureLoggingReq:
+class ConfigureLoggingReq(BaseReq):
     log_requests: Optional[bool] = None
     log_requests_level: Optional[int] = None
     dump_requests_folder: Optional[str] = None
     dump_requests_threshold: Optional[int] = None
+    crash_dump_folder: Optional[str] = None
 
 
 @dataclass
-class OpenSessionReqInput:
+class OpenSessionReqInput(BaseReq):
     capacity_of_str_len: int
     session_id: Optional[str] = None
 
 
 @dataclass
-class CloseSessionReqInput:
+class CloseSessionReqInput(BaseReq):
     session_id: str
 
 
 @dataclass
-class OpenSessionReqOutput:
+class OpenSessionReqOutput(BaseReq):
     session_id: Optional[str]
     success: bool
 
 
 @dataclass
-class HealthCheckOutput:
+class HealthCheckOutput(BaseReq):
     pass
 
 
-class ExpertDistributionReq(Enum):
+class ExpertDistributionReqType(Enum):
     START_RECORD = 1
     STOP_RECORD = 2
     DUMP_RECORD = 3
 
 
+class ExpertDistributionReq(BaseReq):
+    action: ExpertDistributionReqType
+
+
 @dataclass
-class ExpertDistributionReqOutput:
+class ExpertDistributionReqOutput(BaseReq):
     pass
 
 
@@ -1110,7 +1297,7 @@ class Tool:
 
 
 @dataclass
-class ParseFunctionCallReq:
+class ParseFunctionCallReq(BaseReq):
     text: str  # The text to parse.
     tools: List[Tool] = field(
         default_factory=list
@@ -1121,31 +1308,31 @@ class ParseFunctionCallReq:
 
 
 @dataclass
-class SeparateReasoningReqInput:
+class SeparateReasoningReqInput(BaseReq):
     text: str  # The text to parse.
     reasoning_parser: str  # Specify the parser type, e.g., "deepseek-r1".
 
 
 @dataclass
-class VertexGenerateReqInput:
+class VertexGenerateReqInput(BaseReq):
     instances: List[dict]
     parameters: Optional[dict] = None
 
 
 @dataclass
-class RpcReqInput:
+class RpcReqInput(BaseReq):
     method: str
     parameters: Optional[Dict] = None
 
 
 @dataclass
-class RpcReqOutput:
+class RpcReqOutput(BaseReq):
     success: bool
     message: str
 
 
 @dataclass
-class LoadLoRAAdapterReqInput:
+class LoadLoRAAdapterReqInput(BaseReq):
     # The name of the lora module to newly loaded.
     lora_name: str
     # The path of loading.
@@ -1165,7 +1352,7 @@ class LoadLoRAAdapterReqInput:
 
 
 @dataclass
-class UnloadLoRAAdapterReqInput:
+class UnloadLoRAAdapterReqInput(BaseReq):
     # The name of lora module to unload.
     lora_name: str
     # The unique identifier for the LoRA adapter, which automatically generated in the `TokenizerManager`.
@@ -1179,23 +1366,23 @@ class UnloadLoRAAdapterReqInput:
 
 
 @dataclass
-class LoRAUpdateResult:
+class LoRAUpdateOutput(BaseReq):
     success: bool
     error_message: Optional[str] = None
     loaded_adapters: Optional[Dict[str, LoRARef]] = None
 
 
-LoadLoRAAdapterReqOutput = UnloadLoRAAdapterReqOutput = LoRAUpdateResult
+LoadLoRAAdapterReqOutput = UnloadLoRAAdapterReqOutput = LoRAUpdateOutput
 
 
 @dataclass
-class MultiTokenizerRegisterReq:
-    rids: Optional[Union[List[str], str]] = None
+class MultiTokenizerRegisterReq(BaseBatchReq):
     ipc_name: Optional[str] = None
 
 
 @dataclass
-class MultiTokenizerWarpper:
+class MultiTokenizerWrapper:
+    # FIXME(lsyin): remove this
     worker_id: int
     obj: Optional[Any] = None
 
@@ -1206,5 +1393,49 @@ class BlockReqType(Enum):
 
 
 @dataclass
-class BlockReqInput:
+class BlockReqInput(BaseReq):
     type: BlockReqType
+
+
+@dataclass
+class GetLoadReqInput(BaseReq):
+    pass
+
+
+@dataclass
+class GetLoadReqOutput(BaseReq):
+    dp_rank: int
+    num_reqs: int
+    num_waiting_reqs: int
+    num_tokens: int
+
+
+@dataclass
+class WatchLoadUpdateReq(BaseReq):
+    loads: List[GetLoadReqOutput]
+
+
+def _check_all_req_types():
+    """A helper function to check all request types are defined in this file."""
+    import inspect
+    import sys
+
+    all_classes = inspect.getmembers(sys.modules[__name__], inspect.isclass)
+    for class_type in all_classes:
+        # check its name
+        name = class_type[0]
+        is_io_struct = (
+            name.endswith("Req") or name.endswith("Input") or name.endswith("Output")
+        )
+        is_base_req = issubclass(class_type[1], BaseReq) or issubclass(
+            class_type[1], BaseBatchReq
+        )
+        if is_io_struct and not is_base_req:
+            raise ValueError(f"{name} is not a subclass of BaseReq or BaseBatchReq.")
+        if is_base_req and not is_io_struct:
+            raise ValueError(
+                f"{name} is a subclass of BaseReq but not follow the naming convention."
+            )
+
+
+_check_all_req_types()