sglang 0.2.15__py3-none-any.whl → 0.3.1__py3-none-any.whl

sglang/srt/managers/controller_multi.py

@@ -71,12 +71,10 @@ class ControllerMulti:
         self,
         server_args: ServerArgs,
         port_args: PortArgs,
-        model_override_args,
     ):
         # Parse args
         self.server_args = server_args
         self.port_args = port_args
-        self.model_override_args = model_override_args
         self.load_balance_method = LoadBalanceMethod.from_str(
             server_args.load_balance_method
         )
@@ -114,7 +112,6 @@ class ControllerMulti:
                 self.server_args,
                 self.port_args,
                 pipe_controller_writer,
-                self.model_override_args,
                 True,
                 gpu_ids,
                 dp_worker_id,
@@ -189,14 +186,13 @@ def start_controller_process(
     server_args: ServerArgs,
     port_args: PortArgs,
     pipe_writer,
-    model_override_args: dict,
 ):
     """Start a controller process."""
 
     configure_logger(server_args)
 
     try:
-        controller = ControllerMulti(server_args, port_args, model_override_args)
+        controller = ControllerMulti(server_args, port_args)
     except Exception:
         pipe_writer.send(get_exception_traceback())
         raise

sglang/srt/managers/controller_single.py

@@ -40,7 +40,6 @@ class ControllerSingle:
         self,
         server_args: ServerArgs,
         port_args: PortArgs,
-        model_override_args: dict,
         gpu_ids: List[int],
         is_data_parallel_worker: bool,
         dp_worker_id: int,
@@ -76,7 +75,6 @@ class ControllerSingle:
                 tp_rank_range,
                 server_args,
                 port_args.nccl_ports[dp_worker_id],
-                model_override_args,
             )
 
         # Launch tp rank 0
@@ -85,7 +83,6 @@ class ControllerSingle:
             0,
             server_args,
             port_args.nccl_ports[dp_worker_id],
-            model_override_args,
         )
         self.tp_cpu_group = self.tp_server.model_runner.tp_group.cpu_group
 
@@ -126,7 +123,6 @@ def start_controller_process(
     server_args: ServerArgs,
     port_args: PortArgs,
     pipe_writer: multiprocessing.connection.Connection,
-    model_override_args: dict,
     is_data_parallel_worker: bool = False,
     gpu_ids: List[int] = None,
     dp_worker_id: int = None,
@@ -149,7 +145,6 @@ def start_controller_process(
         controller = ControllerSingle(
             server_args,
             port_args,
-            model_override_args,
             gpu_ids,
             is_data_parallel_worker,
             dp_worker_id,

sglang/srt/managers/io_struct.py

@@ -20,7 +20,7 @@ processes (TokenizerManager, DetokenizerManager, Controller).
 
 import copy
 import uuid
-from dataclasses import dataclass, field
+from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
 from sglang.srt.managers.schedule_batch import BaseFinishReason
@@ -43,6 +43,7 @@ class GenerateReqInput:
     # Whether to return logprobs.
     return_logprob: Optional[Union[List[bool], bool]] = None
     # If return logprobs, the start location in the prompt for returning logprobs.
+    # By default, this value is "-1", which means it will only return logprobs for output tokens.
     logprob_start_len: Optional[Union[List[int], int]] = None
     # If return logprobs, the number of top logprobs to return at each position.
     top_logprobs_num: Optional[Union[List[int], int]] = None
@@ -50,6 +51,13 @@ class GenerateReqInput:
     return_text_in_logprobs: bool = False
     # Whether to stream output.
     stream: bool = False
+    # The modalities of the image data [image, multi-images, video]
+    modalities: Optional[List[str]] = None
+
+    is_single: bool = True
+
+    # LoRA related
+    lora_path: Optional[Union[List[Optional[str]], Optional[str]]] = None
 
     def post_init(self):
         if (self.text is None and self.input_ids is None) or (
@@ -177,6 +185,11 @@ class TokenizedGenerateReqInput:
     top_logprobs_num: int
     # Whether to stream output
     stream: bool
+    # Modalities of the input images
+    modalites: Optional[List[str]] = None
+
+    # LoRA related
+    lora_path: Optional[str] = None  # None means just use the base model
 
 
 @dataclass
@@ -190,6 +203,8 @@ class EmbeddingReqInput:
     # Dummy sampling params for compatibility
     sampling_params: Union[List[Dict], Dict] = None
 
+    is_single: bool = True
+
     def post_init(self):
         if (self.text is None and self.input_ids is None) or (
             self.text is not None and self.input_ids is not None

sglang/srt/managers/policy_scheduler.py

@@ -108,18 +108,25 @@ class PrefillAdder:
     def __init__(
         self,
         tree_cache: BasePrefixCache,
+        running_batch: ScheduleBatch,
+        new_token_ratio: float,
         rem_total_tokens: int,
         rem_input_tokens: int,
         rem_chunk_tokens: Optional[int],
         mixed_with_decode_tokens: int = 0,
     ):
         self.tree_cache = tree_cache
+        self.running_batch = running_batch
+        self.new_token_ratio = new_token_ratio
         self.rem_total_tokens = rem_total_tokens - mixed_with_decode_tokens
+        self.rem_total_tokens_ = self.rem_total_tokens
+        self.total_tokens = rem_total_tokens
         self.rem_input_tokens = rem_input_tokens - mixed_with_decode_tokens
         self.rem_chunk_tokens = rem_chunk_tokens
         if self.rem_chunk_tokens is not None:
             self.rem_chunk_tokens -= mixed_with_decode_tokens
 
+        self.req_states = None
         self.can_run_list = []
         self.new_inflight_req = None
         self.log_hit_tokens = 0
@@ -136,16 +143,20 @@ class PrefillAdder:
             )
         )
 
-    def remove_running_tokens(
-        self, running_batch: ScheduleBatch, new_token_ratio: float
-    ):
+    def remove_running_tokens(self, running_batch: ScheduleBatch):
         self.rem_total_tokens -= sum(
             [
                 min(
                     (r.sampling_params.max_new_tokens - len(r.output_ids)),
                     CLIP_MAX_NEW_TOKENS,
                 )
-                * new_token_ratio
+                * self.new_token_ratio
+                for r in running_batch.reqs
+            ]
+        )
+        self.rem_total_tokens_ -= sum(
+            [
+                r.sampling_params.max_new_tokens - len(r.output_ids)
                 for r in running_batch.reqs
             ]
         )
@@ -154,6 +165,7 @@ class PrefillAdder:
         self, prefix_len: int, extend_input_len: int, max_new_tokens: int
     ):
         self.rem_total_tokens -= extend_input_len + max_new_tokens
+        self.rem_total_tokens_ -= extend_input_len + max_new_tokens
         self.rem_input_tokens -= extend_input_len
         if self.rem_chunk_tokens is not None:
             self.rem_chunk_tokens -= extend_input_len
@@ -161,7 +173,29 @@ class PrefillAdder:
         self.log_hit_tokens += prefix_len
         self.log_input_tokens += extend_input_len
 
+    def add_inflight_req_ignore_eos(self, req: Req):
+        truncated = req.extend_input_len > self.rem_chunk_tokens
+        req.extend_input_len = min(req.extend_input_len, self.rem_chunk_tokens)
+        req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
+        self.can_run_list.append(req)
+
+        self._prefill_one_req(
+            0,
+            req.extend_input_len,
+            (
+                min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS)
+                if not truncated
+                else 0
+            ),
+        )
+
+        # Return if chunked prefill not finished
+        return req if truncated else None
+
     def add_inflight_req(self, req: Req):
+        if req.sampling_params.ignore_eos:
+            return self.add_inflight_req_ignore_eos(req)
+
         truncated = req.extend_input_len > self.rem_chunk_tokens
         req.extend_input_len = min(req.extend_input_len, self.rem_chunk_tokens)
         req.fill_ids = req.fill_ids[: len(req.prefix_indices) + req.extend_input_len]
@@ -190,7 +224,90 @@ class PrefillAdder:
             delta = self.tree_cache.dec_lock_ref(last_node)
             self.rem_total_tokens += delta
 
+    def add_one_req_ignore_eos(self, req: Req):
+        def get_req_state(r):
+            new_token_ratio = (
+                1.0 if r.sampling_params.ignore_eos else self.new_token_ratio
+            )
+            tokens_left = r.sampling_params.max_new_tokens * new_token_ratio - len(
+                r.output_ids
+            )
+            tokens_occupied = len(r.origin_input_ids) + len(r.output_ids)
+
+            if tokens_left > 0:
+                return (tokens_left, tokens_occupied)
+
+            return None
+
+        # Quick Check
+        can_run = False
+        if (
+            req.extend_input_len + req.sampling_params.max_new_tokens
+            <= self.rem_total_tokens
+        ):
+            can_run = True
+
+        if not can_run:
+            if self.req_states is None:
+                self.req_states = []
+                if self.running_batch is not None:
+                    for r in self.running_batch.reqs:
+                        state = get_req_state(r)
+                        if state is not None:
+                            self.req_states.append(state)
+                for r in self.can_run_list:
+                    state = get_req_state(r)
+                    if state is not None:
+                        self.req_states.append(state)
+                state = get_req_state(req)
+                if state is not None:
+                    self.req_states.append(state)
+
+                self.req_states.sort(key=lambda x: x[0])
+            else:
+                state = get_req_state(req)
+                if state is not None:
+                    for i, (tokens_left, tokens_occupied) in enumerate(self.req_states):
+                        if tokens_left >= state[0]:
+                            self.req_states.insert(i, state)
+                            break
+                    else:
+                        self.req_states.append(state)
+
+            tokens_freed = 0
+            for i, (tokens_left, tokens_occupied) in enumerate(self.req_states):
+                decode_steps = (
+                    self.req_states[i + 1][0]
+                    if i + 1 < len(self.req_states)
+                    else tokens_left
+                )
+                bs = len(self.req_states) - i
+                if self.total_tokens + tokens_freed - decode_steps * bs <= 0:
+                    return False
+                tokens_freed += tokens_occupied
+
+        if req.extend_input_len <= self.rem_chunk_tokens:
+            self.can_run_list.append(req)
+            self._prefill_one_req(
+                0,
+                req.extend_input_len,
+                min(req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS),
+            )
+        else:
+            # Chunked prefill
+            trunc_len = self.rem_chunk_tokens
+            req.extend_input_len = trunc_len
+            req.fill_ids = req.fill_ids[:trunc_len]
+            self.can_run_list.append(req)
+            self.new_inflight_req = req
+            self._prefill_one_req(0, trunc_len, 0)
+
+        return True
+
     def add_one_req(self, req: Req):
+        if req.sampling_params.ignore_eos and self.tree_cache.disable:
+            return self.add_one_req_ignore_eos(req)
+
         total_tokens = req.extend_input_len + min(
             req.sampling_params.max_new_tokens, CLIP_MAX_NEW_TOKENS
         )
@@ -233,4 +350,4 @@ class PrefillAdder:
                 self.tree_cache.inc_lock_ref(req.last_node)
                 self._prefill_one_req(prefix_len, trunc_len, 0)
 
-        return True
+        return True and not self.no_remaining_tokens()