sglang 0.1.14__py3-none-any.whl → 0.1.16__py3-none-any.whl

sglang/srt/managers/router/scheduler.py

@@ -27,44 +27,33 @@ class Scheduler:
             return forward_queue
         elif self.schedule_heuristic == "fcfs":
             return forward_queue
-        elif self.schedule_heuristic == "weight":
+        elif self.schedule_heuristic == "dfs-weight":
             last_node_to_reqs = defaultdict(list)
             for req in forward_queue:
                 last_node_to_reqs[req.last_node].append(req)
-            for node in last_node_to_reqs:
-                last_node_to_reqs[node].sort(key=lambda x: -len(x.prefix_indices))
 
             node_to_weight = defaultdict(int)
-            self._calc_weight_recursive(
-                self.tree_cache.root_node, last_node_to_reqs, node_to_weight
-            )
+            for node in last_node_to_reqs:
+                node_to_weight[node] = len(last_node_to_reqs[node])
+            self.calc_weight(self.tree_cache.root_node, node_to_weight)
 
-            tmp_queue = []
-            self._get_weight_priority_recursive(
-                self.tree_cache.root_node, node_to_weight, last_node_to_reqs, tmp_queue
+            q = []
+            self.get_dfs_priority(
+                self.tree_cache.root_node, node_to_weight, last_node_to_reqs, q
             )
-            assert len(tmp_queue) == len(forward_queue)
-            return tmp_queue
+            assert len(q) == len(forward_queue)
+            return q
         else:
             raise ValueError(f"Unknown schedule_heuristic: {self.schedule_heuristic}")
 
-    def _calc_weight_recursive(self, cur_node, last_node_to_reqs, node_to_weight):
-        node_to_weight[cur_node] = 1
-        if cur_node in last_node_to_reqs:
-            node_to_weight[cur_node] += len(last_node_to_reqs[cur_node])
+    def calc_weight(self, cur_node, node_to_weight):
         for child in cur_node.children.values():
-            self._calc_weight_recursive(child, last_node_to_reqs, node_to_weight)
+            self.calc_weight(child, node_to_weight)
             node_to_weight[cur_node] += node_to_weight[child]
 
-    def _get_weight_priority_recursive(
-        self, cur_node, node_to_wight, last_node_to_reqs, tmp_queue
-    ):
-        visit_list = [child for child in cur_node.children.values()]
-        visit_list.sort(key=lambda x: -node_to_wight[x])
-        # for node in visit_list:
-        #     print(f"{node_to_wight[node]} {len(node.value) if node.value is not None else 0}")
-        for child in visit_list:
-            self._get_weight_priority_recursive(
-                child, node_to_wight, last_node_to_reqs, tmp_queue
-            )
-        tmp_queue.extend(last_node_to_reqs[cur_node])
+    def get_dfs_priority(self, cur_node, node_to_priority, last_node_to_reqs, q):
+        childs = [child for child in cur_node.children.values()]
+        childs.sort(key=lambda x: -node_to_priority[x])
+        for child in childs:
+            self.get_dfs_priority(child, node_to_priority, last_node_to_reqs, q)
+        q.extend(last_node_to_reqs[cur_node])

sglang/srt/managers/tokenizer_manager.py

@@ -1,6 +1,7 @@
 import asyncio
 import concurrent.futures
 import dataclasses
+import logging
 import multiprocessing as mp
 import os
 from typing import List
@@ -10,6 +11,7 @@ import transformers
 import uvloop
 import zmq
 import zmq.asyncio
+
 from sglang.srt.hf_transformers_utils import (
     get_config,
     get_context_length,
@@ -30,13 +32,14 @@ from sglang.srt.utils import get_exception_traceback, is_multimodal_model, load_
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
+logger = logging.getLogger(__name__)
+
 
 @dataclasses.dataclass
 class ReqState:
     out_list: List
     finished: bool
     event: asyncio.Event
-    lock: asyncio.Lock
 
 
 global global_processor
@@ -57,21 +60,29 @@ def get_pixel_values(
 ):
     try:
         processor = processor or global_processor
-        image = load_image(image_data)
-        image_hash = hash(image_data)
-        if image_aspect_ratio == "pad":
-            image = expand2square(
-                image, tuple(int(x * 255) for x in processor.image_processor.image_mean)
-            )
-            pixel_values = processor.image_processor(image)["pixel_values"][0]
-        elif image_aspect_ratio == "anyres":
-            pixel_values = process_anyres_image(
-                image, processor.image_processor, image_grid_pinpoints
-            )
+        image, image_size = load_image(image_data)
+        if image_size != None:
+            image_hash = hash(image_data)
+            pixel_values = processor.image_processor(image)["pixel_values"]
+            for _ in range(len(pixel_values)):
+                pixel_values[_] = pixel_values[_].astype(np.float16)
+            pixel_values = np.stack(pixel_values, axis=0)
+            return pixel_values, image_hash, image_size
         else:
-            pixel_values = processor.image_processor(image)["pixel_values"][0]
-        pixel_values = pixel_values.astype(np.float16)
-        return pixel_values, image_hash, image.size
+            image_hash = hash(image_data)
+            if image_aspect_ratio == "pad":
+                image = expand2square(
+                    image, tuple(int(x * 255) for x in processor.image_processor.image_mean)
+                )
+                pixel_values = processor.image_processor(image)["pixel_values"][0]
+            elif image_aspect_ratio == "anyres":
+                pixel_values = process_anyres_image(
+                    image, processor.image_processor, image_grid_pinpoints
+                )
+            else:
+                pixel_values = processor.image_processor(image)["pixel_values"][0]
+            pixel_values = pixel_values.astype(np.float16)
+            return pixel_values, image_hash, image.size
     except Exception:
         print("Exception in TokenizerManager:\n" + get_exception_traceback())
 
@@ -81,6 +92,7 @@ class TokenizerManager:
         self,
         server_args: ServerArgs,
         port_args: PortArgs,
+        model_overide_args: dict = None,
     ):
         self.server_args = server_args
 
@@ -93,9 +105,10 @@ class TokenizerManager:
 
         self.model_path = server_args.model_path
         self.hf_config = get_config(
-            self.model_path, trust_remote_code=server_args.trust_remote_code
+            self.model_path,
+            trust_remote_code=server_args.trust_remote_code,
+            model_overide_args=model_overide_args,
         )
-
         self.context_len = get_context_length(self.hf_config)
 
         if is_multimodal_model(self.model_path):
@@ -144,11 +157,21 @@ class TokenizerManager:
         if self.to_create_loop:
             await self.create_handle_loop()
 
-        is_single = isinstance(obj.text, str)
-
+        is_single = obj.is_single
         if is_single:
             rid = obj.rid
-            input_ids = self.tokenizer.encode(obj.text)
+
+            if obj.input_ids is None:
+                input_ids = self.tokenizer.encode(obj.text)
+            else:
+                input_ids = obj.input_ids
+
+            if len(input_ids) >= self.context_len:
+                raise ValueError(
+                    f"The input ({len(input_ids)} tokens) is longer than the "
+                    f"model's context length ({self.context_len} tokens)"
+                )
+
             sampling_params = SamplingParams(**obj.sampling_params)
             if sampling_params.max_new_tokens != 0:
                 sampling_params.normalize(self.tokenizer)
@@ -174,18 +197,26 @@ class TokenizerManager:
                 sampling_params=sampling_params,
                 return_logprob=obj.return_logprob,
                 logprob_start_len=obj.logprob_start_len,
+                top_logprobs_num=obj.top_logprobs_num,
                 stream=obj.stream,
             )
             self.send_to_router.send_pyobj(tokenized_obj)
 
-            lock = asyncio.Lock()
             event = asyncio.Event()
-            state = ReqState([], False, event, lock)
+            state = ReqState([], False, event)
             self.rid_to_state[rid] = state
 
             while True:
                 await event.wait()
-                yield state.out_list[-1]
+                out = self.convert_logprob_style(state.out_list[-1],
+                                                 obj.return_logprob,
+                                                 obj.top_logprobs_num,
+                                                 obj.return_text_in_logprobs)
+
+                if self.server_args.log_requests and state.finished:
+                    logger.info(f"in={obj.text}, out={out}")
+
+                yield out
                 state.out_list = []
                 if state.finished:
                     del self.rid_to_state[rid]
@@ -193,10 +224,22 @@ class TokenizerManager:
                 event.clear()
         else:
             assert obj.stream is False
-            bs = len(obj.text)
+
+            if obj.input_ids is None:
+                bs = len(obj.text)
+            else:
+                bs = len(obj.input_ids)
+
             for i in range(bs):
                 rid = obj.rid[i]
-                input_ids = self.tokenizer.encode(obj.text[i])
+
+                if obj.input_ids is None:
+                    input_text = obj.text[i]
+                    input_ids = self.tokenizer.encode(obj.text[i])
+                else:
+                    input_text = None
+                    input_ids = obj.input_ids[i]
+
                 sampling_params = SamplingParams(**obj.sampling_params[i])
                 if sampling_params.max_new_tokens != 0:
                     sampling_params.normalize(self.tokenizer)
@@ -209,7 +252,7 @@ class TokenizerManager:
                     )
                 tokenized_obj = TokenizedGenerateReqInput(
                     rid=rid,
-                    input_text=obj.text[i],
+                    input_text=input_text,
                     input_ids=input_ids,
                     pixel_values=pixel_values,
                     image_hash=image_hash,
@@ -217,13 +260,13 @@ class TokenizerManager:
                     sampling_params=sampling_params,
                     return_logprob=obj.return_logprob[i],
                     logprob_start_len=obj.logprob_start_len[i],
+                    top_logprobs_num=obj.top_logprobs_num[i],
                     stream=obj.stream,
                 )
                 self.send_to_router.send_pyobj(tokenized_obj)
 
-                lock = asyncio.Lock()
                 event = asyncio.Event()
-                state = ReqState([], False, event, lock)
+                state = ReqState([], False, event)
                 self.rid_to_state[rid] = state
 
             output_list = []
@@ -231,16 +274,16 @@ class TokenizerManager:
                 rid = obj.rid[i]
                 state = self.rid_to_state[rid]
                 await state.event.wait()
-                output_list.append(state.out_list[-1])
+                output_list.append(
+                    self.convert_logprob_style(state.out_list[-1],
+                                               obj.return_logprob[i],
+                                               obj.top_logprobs_num[i],
+                                               obj.return_text_in_logprobs))
                 assert state.finished
                 del self.rid_to_state[rid]
 
             yield output_list
 
-    async def detokenize(self, obj: DetokenizeReqInput):
-        token_texts = self.tokenizer.convert_ids_to_tokens(obj.input_ids)
-        return [t.decode() if isinstance(t, bytes) else t for t in token_texts]
-
     async def flush_cache(self):
         flush_cache_req = FlushCacheReq()
         self.send_to_router.send_pyobj(flush_cache_req)
@@ -267,3 +310,37 @@ class TokenizerManager:
                     state.event.set()
             else:
                 raise ValueError(f"Invalid object: {recv_obj}")
+
+    def convert_logprob_style(self, ret, return_logprob, top_logprobs_num, return_text_in_logprobs):
+        if return_logprob:
+            ret["meta_info"]["prefill_token_logprobs"] = self.detokenize_logprob_tokens(
+                ret["meta_info"]["prefill_token_logprobs"], return_text_in_logprobs
+            )
+            ret["meta_info"]["decode_token_logprobs"] = self.detokenize_logprob_tokens(
+                ret["meta_info"]["decode_token_logprobs"], return_text_in_logprobs
+            )
+        if top_logprobs_num > 0:
+            ret["meta_info"]["prefill_top_logprobs"] = self.detokenize_top_logprobs_tokens(
+                ret["meta_info"]["prefill_top_logprobs"], return_text_in_logprobs
+            )
+            ret["meta_info"]["decode_top_logprobs"] = self.detokenize_top_logprobs_tokens(
+                ret["meta_info"]["decode_top_logprobs"], return_text_in_logprobs
+            )
+        return ret
+
+    def detokenize_logprob_tokens(self, token_logprobs, decode_to_text):
+        if not decode_to_text:
+            return [(logprob, token_id, None) for logprob, token_id in token_logprobs]
+
+        token_ids = [tid for _, tid in token_logprobs]
+        token_texts = self.tokenizer.batch_decode(token_ids)
+        return [
+            (logprob, token_id, token_text)
+            for (logprob, token_id), token_text, in zip(token_logprobs, token_texts)
+        ]
+
+    def detokenize_top_logprobs_tokens(self, top_logprobs, decode_to_text):
+        for i, t in enumerate(top_logprobs):
+            if t:
+                top_logprobs[i] = self.detokenize_logprob_tokens(t, decode_to_text)
+        return top_logprobs

sglang/srt/memory_pool.py

@@ -31,9 +31,6 @@ class ReqToTokenPool:
             self.can_use_mem_size += free_index.shape[0]
         self.mem_state[free_index] = 1
 
-        # if self.can_use_mem_size == len(self.mem_state):
-        #     print(f"ReqToTokenPool: freed all. size = {self.can_use_mem_size}.")
-
     def clear(self):
         self.mem_state.fill_(1)
         self.can_use_mem_size = len(self.mem_state)
@@ -42,7 +39,7 @@ class ReqToTokenPool:
 class TokenToKVPool:
     def __init__(self, size, dtype, head_num, head_dim, layer_num):
         self.mem_state = torch.zeros((size,), dtype=torch.int16, device="cuda")
-        self.alloc_ct = 0
+        self.total_ref_ct = 0
 
         # [size, key/value, head_num, head_dim] for each layer
         self.kv_data = [
@@ -83,9 +80,6 @@ class TokenToKVPool:
         self.add_refs(select_index)
         return select_index.to(torch.int32), start_loc, start_loc + need_size
 
-    def free(self, free_index):
-        return self.decrease_refs(free_index)
-
     def used_size(self):
         return len(torch.nonzero(self.mem_state).squeeze(1))
 
@@ -93,20 +87,17 @@ class TokenToKVPool:
         return torch.sum(self.mem_state == 0).item()
 
     def add_refs(self, token_index: torch.Tensor):
-        self.alloc_ct += len(token_index)
+        self.total_ref_ct += len(token_index)
         self.mem_state[token_index] += 1
 
-    def decrease_refs(self, token_index: torch.Tensor):
-        self.alloc_ct -= len(token_index)
+    def dec_refs(self, token_index: torch.Tensor):
+        self.total_ref_ct -= len(token_index)
         self.mem_state[token_index] -= 1
 
         num_freed = torch.sum(self.mem_state[token_index] == 0)
 
-        # if self.alloc_ct == 0:
-        #     print(f"TokenToKVPool: freed all. size = {len(self.mem_state)}.")
-
         return num_freed
 
     def clear(self):
         self.mem_state.fill_(0)
-        self.alloc_ct = 0
+        self.total_ref_ct = 0

sglang/srt/model_config.py

@@ -10,12 +10,16 @@ class ModelConfig:
         trust_remote_code: bool = True,
         revision: Optional[str] = None,
         context_length: Optional[int] = None,
+        model_overide_args: Optional[dict] = None,
     ) -> None:
         self.path = path
         self.trust_remote_code = trust_remote_code
         self.revision = revision
         self.hf_config = get_config(self.path, trust_remote_code, revision)
 
+        if model_overide_args is not None:
+            self.hf_config.update(model_overide_args)
+
         if context_length is not None:
             self.context_len = context_length
         else:
@@ -29,6 +33,13 @@ class ModelConfig:
         )
         self.num_attention_heads = self.hf_config.num_attention_heads
         self.num_key_value_heads = getattr(self.hf_config, "num_key_value_heads", None)
+
+        # for Dbrx and MPT models
+        if self.hf_config.model_type in ["dbrx", "mpt"]:
+            self.num_key_value_heads = getattr(
+                self.hf_config.attn_config, "kv_n_heads", None
+            )
+
         if self.num_key_value_heads is None:
             self.num_key_value_heads = self.num_attention_heads
         self.hidden_size = self.hf_config.hidden_size