sglang 0.2.4__py3-none-any.whl → 0.2.6__py3-none-any.whl

sglang/srt/managers/controller/tp_worker.py

@@ -98,17 +98,21 @@ class ModelTpServer:
             if server_args.max_prefill_tokens is None
             else server_args.max_prefill_tokens
         )
-        self.max_running_requests = (
-            self.max_total_num_tokens // 2
-            if server_args.max_running_requests is None
-            else server_args.max_running_requests
-        )
         self.max_running_requests = min(
-            self.max_running_requests, self.model_runner.req_to_token_pool.size - 1
+            (
+                self.max_total_num_tokens // 2
+                if server_args.max_running_requests is None
+                else server_args.max_running_requests
+            ),
+            self.model_runner.req_to_token_pool.size - 1,
         )
         self.int_token_logit_bias = torch.tensor(
             get_int_token_logit_bias(self.tokenizer, self.model_config.vocab_size)
         )
+        self.max_req_input_len = min(
+            self.model_config.context_len - 1,
+            self.max_total_num_tokens - 1,
+        )
         set_random_seed(server_args.random_seed)
 
         # Print info
@@ -295,18 +299,20 @@ class ModelTpServer:
                 )
 
         # Truncate prompts that are too long
-        req.origin_input_ids = req.origin_input_ids[: self.model_config.context_len - 1]
+        if len(req.origin_input_ids) >= self.max_req_input_len:
+            logger.warn(
+                "Request length is longer than the KV cache pool size or "
+                "the max context length. Truncated!!!"
+            )
+            req.origin_input_ids = req.origin_input_ids[: self.max_req_input_len]
         req.sampling_params.max_new_tokens = min(
-            req.sampling_params.max_new_tokens,
-            self.model_config.context_len - 1 - len(req.origin_input_ids),
-            self.max_total_num_tokens - 128 - len(req.origin_input_ids),
+            (
+                req.sampling_params.max_new_tokens
+                if req.sampling_params.max_new_tokens is not None
+                else 1 << 30
+            ),
+            self.max_req_input_len - 1 - len(req.origin_input_ids),
         )
-        if req.sampling_params.max_new_tokens < 0:
-            req.origin_input_ids = req.origin_input_ids[
-                : self.max_total_num_tokens - 128
-            ]
-            logger.error("Request longer than memory pool size, truncated!!!")
-
         self.forward_queue.append(req)
 
     def get_new_prefill_batch(self) -> Optional[Batch]:
@@ -449,7 +455,7 @@ class ModelTpServer:
                     torch.arange(len(next_token_ids), device=next_token_ids.device),
                     next_token_ids,
                 ].tolist()
-                output.prefill_token_logprobs = output.prefill_token_logprobs.tolist()
+                output.input_token_logprobs = output.input_token_logprobs.tolist()
                 output.normalized_prompt_logprobs = (
                     output.normalized_prompt_logprobs.tolist()
                 )
@@ -475,24 +481,24 @@ class ModelTpServer:
         if req.normalized_prompt_logprob is None:
             req.normalized_prompt_logprob = output.normalized_prompt_logprobs[i]
 
-        if req.prefill_token_logprobs is None:
+        if req.input_token_logprobs is None:
             # If logprob_start_len > 0, then first logprob_start_len prompt tokens will be ignored.
-            req.prefill_token_logprobs = list(
+            req.input_token_logprobs = list(
                 zip(
-                    output.prefill_token_logprobs[pt : pt + req.extend_input_len - 1],
+                    output.input_token_logprobs[pt : pt + req.extend_input_len - 1],
                     req.input_ids[-req.extend_input_len + 1 :],
                 )
             )
             if req.logprob_start_len == 0:
-                req.prefill_token_logprobs = [
+                req.input_token_logprobs = [
                     (None, req.input_ids[0])
-                ] + req.prefill_token_logprobs
+                ] + req.input_token_logprobs
 
         if req.last_update_decode_tokens != 0:
-            req.decode_token_logprobs.extend(
+            req.output_token_logprobs.extend(
                 list(
                     zip(
-                        output.prefill_token_logprobs[
+                        output.input_token_logprobs[
                             pt
                             + req.extend_input_len
                             - req.last_update_decode_tokens : pt
@@ -504,21 +510,21 @@ class ModelTpServer:
                 )
             )
 
-        req.decode_token_logprobs.append(
+        req.output_token_logprobs.append(
             (output.next_token_logprobs[i], next_token_ids[i])
         )
 
         if req.top_logprobs_num > 0:
-            if req.prefill_top_logprobs is None:
-                req.prefill_top_logprobs = output.prefill_top_logprobs[i]
+            if req.input_top_logprobs is None:
+                req.input_top_logprobs = output.input_top_logprobs[i]
                 if req.logprob_start_len == 0:
-                    req.prefill_top_logprobs = [None] + req.prefill_top_logprobs
+                    req.input_top_logprobs = [None] + req.input_top_logprobs
 
             if req.last_update_decode_tokens != 0:
-                req.decode_top_logprobs.extend(
-                    output.prefill_top_logprobs[i][-req.last_update_decode_tokens + 1 :]
+                req.output_top_logprobs.extend(
+                    output.input_top_logprobs[i][-req.last_update_decode_tokens + 1 :]
                 )
-            req.decode_top_logprobs.append(output.decode_top_logprobs[i])
+            req.output_top_logprobs.append(output.output_top_logprobs[i])
 
     def cache_filled_batch(self, batch: Batch):
         req_pool_indices_cpu = batch.req_pool_indices.cpu().numpy()
@@ -583,11 +589,11 @@ class ModelTpServer:
             req.check_finished()
 
             if req.return_logprob:
-                req.decode_token_logprobs.append(
+                req.output_token_logprobs.append(
                     (next_token_logprobs[i], next_token_id)
                 )
                 if req.top_logprobs_num > 0:
-                    req.decode_top_logprobs.append(output.decode_top_logprobs[i])
+                    req.output_top_logprobs.append(output.output_top_logprobs[i])
 
         self.handle_finished_requests(batch)
 
@@ -639,16 +645,16 @@ class ModelTpServer:
                 }
                 if req.return_logprob:
                     (
-                        meta_info["prefill_token_logprobs"],
-                        meta_info["decode_token_logprobs"],
-                        meta_info["prefill_top_logprobs"],
-                        meta_info["decode_top_logprobs"],
+                        meta_info["input_token_logprobs"],
+                        meta_info["output_token_logprobs"],
+                        meta_info["input_top_logprobs"],
+                        meta_info["output_top_logprobs"],
                         meta_info["normalized_prompt_logprob"],
                     ) = (
-                        req.prefill_token_logprobs,
-                        req.decode_token_logprobs,
-                        req.prefill_top_logprobs,
-                        req.decode_top_logprobs,
+                        req.input_token_logprobs,
+                        req.output_token_logprobs,
+                        req.input_top_logprobs,
+                        req.output_top_logprobs,
                         req.normalized_prompt_logprob,
                     )
                 output_meta_info.append(meta_info)

sglang/srt/managers/io_struct.py

@@ -20,7 +20,7 @@ class GenerateReqInput:
     # The image input. It can be a file name, a url, or base64 encoded string.
     # See also python/sglang/srt/utils.py:load_image.
     image_data: Optional[Union[List[str], str]] = None
-    # The sampling_params.
+    # The sampling_params. See descriptions below.
     sampling_params: Union[List[Dict], Dict] = None
     # The request id.
     rid: Optional[Union[List[str], str]] = None
@@ -30,7 +30,7 @@ class GenerateReqInput:
     logprob_start_len: Optional[Union[List[int], int]] = None
     # The number of top logprobs to return.
     top_logprobs_num: Optional[Union[List[int], int]] = None
-    # Whether to detokenize tokens in logprobs.
+    # Whether to detokenize tokens in text in the returned logprobs.
     return_text_in_logprobs: bool = False
     # Whether to stream output.
     stream: bool = False

sglang/srt/managers/tokenizer_manager.py

@@ -133,24 +133,10 @@ class TokenizerManager:
             async for response in self._handle_batch_request(obj, request):
                 yield response
 
-    async def _handle_single_request(self, obj, request, index=None, is_prefill=False):
-        if is_prefill:
-            if isinstance(obj.text, list):
-                input_text = obj.text[index]
-                rid = obj.rid[index]
-            else:
-                input_text = obj.text
-                rid = obj.rid[0]
-            input_ids = self.tokenizer.encode(input_text)
-            sampling_params = SamplingParams(**obj.sampling_params[0])
-            sampling_params.max_new_tokens = 0
-            pixel_values, image_hash, image_size = await self._get_pixel_values(
-                obj.image_data[0]
-            )
-            return_logprob = obj.return_logprob[0]
-            logprob_start_len = obj.logprob_start_len[0]
-            top_logprobs_num = obj.top_logprobs_num[0]
-        else:
+    async def _handle_single_request(
+        self, obj, request, index=None, is_cache_for_prefill=False
+    ):
+        if not is_cache_for_prefill:
             rid = obj.rid if index is None else obj.rid[index]
             input_text = obj.text if index is None else obj.text[index]
             input_ids = (
@@ -177,6 +163,22 @@ class TokenizerManager:
             top_logprobs_num = (
                 obj.top_logprobs_num if index is None else obj.top_logprobs_num[index]
             )
+        else:
+            if isinstance(obj.text, list):
+                input_text = obj.text[index]
+                rid = obj.rid[index]
+            else:
+                input_text = obj.text
+                rid = obj.rid[0]
+            input_ids = self.tokenizer.encode(input_text)
+            sampling_params = SamplingParams(**obj.sampling_params[0])
+            sampling_params.max_new_tokens = 0
+            pixel_values, image_hash, image_size = await self._get_pixel_values(
+                obj.image_data[0]
+            )
+            return_logprob = obj.return_logprob[0]
+            logprob_start_len = obj.logprob_start_len[0]
+            top_logprobs_num = obj.top_logprobs_num[0]
 
         tokenized_obj = TokenizedGenerateReqInput(
             rid,
@@ -196,26 +198,26 @@ class TokenizerManager:
         event = asyncio.Event()
         state = ReqState([], False, event)
         self.rid_to_state[rid] = state
-        if is_prefill:
-            await self._wait_for_prefill_response(event, state, obj, request, rid)
-            yield input_ids
-        else:
+        if not is_cache_for_prefill:
             async for response in self._wait_for_response(
                 event, state, obj, rid, request
             ):
                 yield response
+        else:
+            await self._wait_for_cache_prefill_response(event, state, obj, rid, request)
+            yield input_ids
 
-    async def _handle_batch_request(self, obj, request):
+    async def _handle_batch_request(self, obj: GenerateReqInput, request):
         batch_size = obj.batch_size
         parallel_sample_num = obj.sampling_params[0].get("n", 1)
 
         if parallel_sample_num != 1:
-            ## send prefill requests
+            # Send prefill requests to cache the common input
             parallel_sample_num += 1
             input_id_result = [] if obj.input_ids is None else None
             for i in range(batch_size):
                 async for input_id in self._handle_single_request(
-                    obj, request, index=i, is_prefill=True
+                    obj, request, index=i, is_cache_for_prefill=True
                 ):
                     if input_id_result is not None:
                         input_id_result.append(input_id)
@@ -224,6 +226,7 @@ class TokenizerManager:
                 obj.input_ids = input_id_result
             elif input_id_result is not None:
                 obj.input_ids = input_id_result[0]
+
         # First send out all requests
         for i in range(batch_size):
             for j in range(parallel_sample_num):
@@ -308,17 +311,15 @@ class TokenizerManager:
 
         yield output_list
 
-    def _validate_input_length(self, input_ids):
+    def _validate_input_length(self, input_ids: List[int]):
         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)."
             )
 
-    def _get_sampling_params(self, sampling_params_data, max_new_tokens=None):
+    def _get_sampling_params(self, sampling_params_data: dict):
         sampling_params = SamplingParams(**sampling_params_data)
-        if max_new_tokens is not None:
-            sampling_params.max_new_tokens = max_new_tokens
         if sampling_params.max_new_tokens != 0:
             sampling_params.normalize(self.tokenizer)
             sampling_params.verify()
@@ -332,7 +333,14 @@ class TokenizerManager:
         else:
             return None, None, None
 
-    async def _wait_for_response(self, event, state, obj, rid, request):
+    async def _wait_for_response(
+        self,
+        event: asyncio.Event,
+        state: ReqState,
+        obj: GenerateReqInput,
+        rid: str,
+        request,
+    ):
         while True:
             try:
                 await asyncio.wait_for(event.wait(), timeout=4)
@@ -361,7 +369,14 @@ class TokenizerManager:
             event.clear()
             yield out
 
-    async def _wait_for_prefill_response(self, event, state, obj, request, rid):
+    async def _wait_for_cache_prefill_response(
+        self,
+        event: asyncio.Event,
+        state: ReqState,
+        obj: GenerateReqInput,
+        rid: str,
+        request,
+    ):
         while True:
             try:
                 await asyncio.wait_for(state.event.wait(), timeout=4)
@@ -380,7 +395,7 @@ class TokenizerManager:
         req = FlushCacheReq()
         self.send_to_router.send_pyobj(req)
 
-    def abort_request(self, rid):
+    def abort_request(self, rid: str):
         if rid not in self.rid_to_state:
             return
         del self.rid_to_state[rid]
@@ -426,31 +441,35 @@ class TokenizerManager:
                 state.event.set()
 
     def convert_logprob_style(
-        self, ret, return_logprob, top_logprobs_num, return_text_in_logprobs
+        self,
+        ret: dict,
+        return_logprob: bool,
+        top_logprobs_num: int,
+        return_text_in_logprobs: bool,
     ):
         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"]["input_token_logprobs"] = self.detokenize_logprob_tokens(
+                ret["meta_info"]["input_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
+            ret["meta_info"]["output_token_logprobs"] = self.detokenize_logprob_tokens(
+                ret["meta_info"]["output_token_logprobs"], return_text_in_logprobs
             )
 
             if top_logprobs_num > 0:
-                ret["meta_info"]["prefill_top_logprobs"] = (
+                ret["meta_info"]["input_top_logprobs"] = (
                     self.detokenize_top_logprobs_tokens(
-                        ret["meta_info"]["prefill_top_logprobs"],
+                        ret["meta_info"]["input_top_logprobs"],
                         return_text_in_logprobs,
                     )
                 )
-                ret["meta_info"]["decode_top_logprobs"] = (
+                ret["meta_info"]["output_top_logprobs"] = (
                     self.detokenize_top_logprobs_tokens(
-                        ret["meta_info"]["decode_top_logprobs"], return_text_in_logprobs
+                        ret["meta_info"]["output_top_logprobs"], return_text_in_logprobs
                     )
                 )
         return ret
 
-    def detokenize_logprob_tokens(self, token_logprobs, decode_to_text):
+    def detokenize_logprob_tokens(self, token_logprobs, decode_to_text: bool):
         if not decode_to_text:
             return [(logprob, token_id, None) for logprob, token_id in token_logprobs]
 
@@ -461,7 +480,7 @@ class TokenizerManager:
             for (logprob, token_id), token_text, in zip(token_logprobs, token_texts)
         ]
 
-    def detokenize_top_logprobs_tokens(self, top_logprobs, decode_to_text):
+    def detokenize_top_logprobs_tokens(self, top_logprobs, decode_to_text: bool):
         for i, t in enumerate(top_logprobs):
             if t:
                 top_logprobs[i] = self.detokenize_logprob_tokens(t, decode_to_text)

sglang/srt/model_config.py

@@ -36,6 +36,11 @@ class ModelConfig:
             "head_dim",
             self.hf_config.hidden_size // self.hf_config.num_attention_heads,
         )
+
+        # FIXME: temporary special judge for deepseek v2 MLA architecture
+        if "DeepseekV2ForCausalLM" in self.hf_config.architectures:
+            self.head_dim = 256
+
         self.num_attention_heads = self.hf_config.num_attention_heads
         self.num_key_value_heads = getattr(self.hf_config, "num_key_value_heads", None)