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

sglang/srt/managers/detokenizer_manager.py

@@ -1,13 +1,17 @@
+"""DetokenizerManager is a process that detokenizes the token ids."""
+
 import asyncio
+import inspect
 
 import uvloop
 import zmq
 import zmq.asyncio
 
 from sglang.srt.hf_transformers_utils import get_tokenizer
+from sglang.srt.managers.controller.infer_batch import FINISH_MATCHED_STR
 from sglang.srt.managers.io_struct import BatchStrOut, BatchTokenIDOut
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import get_exception_traceback
+from sglang.utils import get_exception_traceback, graceful_registry
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
@@ -33,51 +37,41 @@ class DetokenizerManager:
 
     async def handle_loop(self):
         while True:
-            recv_obj = await self.recv_from_router.recv_pyobj()
-
-            if isinstance(recv_obj, BatchTokenIDOut):
-                output_tokens = recv_obj.output_tokens
-
-                # TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request
-                output_strs = self.tokenizer.batch_decode(
-                    output_tokens,
-                    skip_special_tokens=recv_obj.skip_special_tokens[0],
-                    spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[
-                        0
-                    ],
-                )
-
-                # Trim stop str
-                # TODO(lmzheng): handle the case where multiple stop strs are hit
-                for i in range(len(output_strs)):
-                    if recv_obj.hit_stop_str[i] is not None:
-                        pos = output_strs[i].find(recv_obj.hit_stop_str[i])
-                        if pos != -1:
-                            output_strs[i] = output_strs[i][:pos]
-
-                    if len(output_tokens[i]) > 0:
-                        first_token = self.tokenizer.convert_ids_to_tokens(
-                            int(output_tokens[i][0])
-                        )
-                        if not isinstance(first_token, str):
-                            first_token = first_token.decode("utf-8", errors="ignore")
-                        if first_token.startswith("▁"):
-                            output_strs[i] = " " + output_strs[i]
-
-                    output_strs[i] = (
-                        recv_obj.output_and_jump_forward_strs[i] + output_strs[i]
-                    )
-
-                self.send_to_tokenizer.send_pyobj(
-                    BatchStrOut(
-                        recv_obj.rids,
-                        output_strs,
-                        recv_obj.meta_info,
-                        recv_obj.finished,
-                    )
+            recv_obj: BatchTokenIDOut = await self.recv_from_router.recv_pyobj()
+            assert isinstance(recv_obj, BatchTokenIDOut)
+
+            # TODO(lmzheng): handle skip_special_tokens/spaces_between_special_tokens per request
+            surr_texts = self.tokenizer.batch_decode(
+                recv_obj.surr_output_ids,
+                skip_special_tokens=recv_obj.skip_special_tokens[0],
+                spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
+            )
+            read_texts = self.tokenizer.batch_decode(
+                recv_obj.read_output_ids,
+                skip_special_tokens=recv_obj.skip_special_tokens[0],
+                spaces_between_special_tokens=recv_obj.spaces_between_special_tokens[0],
+            )
+
+            # Trim stop str
+            # TODO(lmzheng): handle the case where multiple stop strs are hit
+            output_strs = []
+            for i in range(len(recv_obj.rids)):
+                new_text = read_texts[i][len(surr_texts[i]) :]
+                output_strs.append(recv_obj.decoded_texts[i] + new_text)
+
+                if isinstance(recv_obj.finished_reason[i], FINISH_MATCHED_STR):
+                    pos = output_strs[i].find(recv_obj.finished_reason[i].matched)
+                    if pos != -1:
+                        output_strs[i] = output_strs[i][:pos]
+
+            self.send_to_tokenizer.send_pyobj(
+                BatchStrOut(
+                    rids=recv_obj.rids,
+                    output_str=output_strs,
+                    meta_info=recv_obj.meta_info,
+                    finished_reason=recv_obj.finished_reason,
                 )
-            else:
-                raise ValueError(f"Invalid object: {recv_obj}")
+            )
 
 
 def start_detokenizer_process(
@@ -85,9 +79,11 @@ def start_detokenizer_process(
     port_args: PortArgs,
     pipe_writer,
 ):
+    graceful_registry(inspect.currentframe().f_code.co_name)
+
     try:
         manager = DetokenizerManager(server_args, port_args)
-    except Exception as e:
+    except Exception:
         pipe_writer.send(get_exception_traceback())
         raise
     pipe_writer.send("init ok")

sglang/srt/managers/io_struct.py

@@ -1,7 +1,13 @@
+"""
+The definition of objects transfered between different
+processes (TokenizerManager, DetokenizerManager, Controller).
+"""
+
 import uuid
 from dataclasses import dataclass
 from typing import Dict, List, Optional, Union
 
+from sglang.srt.managers.controller.infer_batch import BaseFinishReason
 from sglang.srt.sampling_params import SamplingParams
 
 
@@ -27,14 +33,12 @@ class GenerateReqInput:
     return_text_in_logprobs: bool = False
     # Whether to stream output
     stream: bool = False
-    # TODO: make all parameters a Union[List[T], T] to allow for batched requests
 
     def post_init(self):
-
-        if self.text is None:
-            assert self.input_ids is not None, "Either text or input_ids should be provided"
-        else:
-            assert self.input_ids is None, "Either text or input_ids should be provided"
+        if (self.text is None and self.input_ids is None) or (
+            self.text is not None and self.input_ids is not None
+        ):
+            raise ValueError("Either text or input_ids should be provided.")
 
         if self.text is not None:
             is_single = isinstance(self.text, str)
@@ -69,7 +73,8 @@ class GenerateReqInput:
             if self.rid is None:
                 self.rid = [uuid.uuid4().hex for _ in range(num)]
             else:
-                assert isinstance(self.rid, list)
+                if not isinstance(self.rid, list):
+                    raise ValueError("The rid should be a list.")
 
             if self.return_logprob is None:
                 self.return_logprob = [False] * num
@@ -105,13 +110,13 @@ class TokenizedGenerateReqInput:
 @dataclass
 class BatchTokenIDOut:
     rids: List[str]
-    output_tokens: List[List[int]]
-    output_and_jump_forward_strs: List[str]
-    hit_stop_str: List[Optional[str]]
+    decoded_texts: List[str]
+    surr_output_ids: List[List[int]]
+    read_output_ids: List[List[int]]
     skip_special_tokens: List[bool]
     spaces_between_special_tokens: List[bool]
     meta_info: List[Dict]
-    finished: List[bool]
+    finished_reason: List[BaseFinishReason]
 
 
 @dataclass
@@ -119,7 +124,7 @@ class BatchStrOut:
     rids: List[str]
     output_str: List[str]
     meta_info: List[Dict]
-    finished: List[bool]
+    finished_reason: List[BaseFinishReason]
 
 
 @dataclass
@@ -127,6 +132,11 @@ class FlushCacheReq:
     pass
 
 
+@dataclass
+class AbortReq:
+    rid: str
+
+
 @dataclass
 class DetokenizeReqInput:
     input_ids: List[int]

sglang/srt/managers/tokenizer_manager.py

@@ -1,16 +1,19 @@
+"""TokenizerManager is a process that tokenizes the text."""
+
 import asyncio
 import concurrent.futures
 import dataclasses
 import logging
 import multiprocessing as mp
 import os
-from typing import List
+from typing import Dict, List
 
 import numpy as np
 import transformers
 import uvloop
 import zmq
 import zmq.asyncio
+from fastapi import BackgroundTasks
 
 from sglang.srt.hf_transformers_utils import (
     get_config,
@@ -19,8 +22,9 @@ from sglang.srt.hf_transformers_utils import (
     get_tokenizer,
 )
 from sglang.srt.managers.io_struct import (
+    AbortReq,
     BatchStrOut,
-    DetokenizeReqInput,
+    BatchTokenIDOut,
     FlushCacheReq,
     GenerateReqInput,
     TokenizedGenerateReqInput,
@@ -28,7 +32,8 @@ from sglang.srt.managers.io_struct import (
 from sglang.srt.mm_utils import expand2square, process_anyres_image
 from sglang.srt.sampling_params import SamplingParams
 from sglang.srt.server_args import PortArgs, ServerArgs
-from sglang.srt.utils import get_exception_traceback, is_multimodal_model, load_image
+from sglang.srt.utils import is_multimodal_model, load_image
+from sglang.utils import get_exception_traceback
 
 asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
 
@@ -42,51 +47,6 @@ class ReqState:
     event: asyncio.Event
 
 
-global global_processor
-
-
-def init_global_processor(server_args: ServerArgs):
-    global global_processor
-    transformers.logging.set_verbosity_error()
-    global_processor = get_processor(
-        server_args.tokenizer_path,
-        tokenizer_mode=server_args.tokenizer_mode,
-        trust_remote_code=server_args.trust_remote_code,
-    )
-
-
-def get_pixel_values(
-    image_data, image_aspect_ratio=None, image_grid_pinpoints=None, processor=None
-):
-    try:
-        processor = processor or global_processor
-        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:
-            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())
-
-
 class TokenizerManager:
     def __init__(
         self,
@@ -132,7 +92,7 @@ class TokenizerManager:
             )
 
         self.to_create_loop = True
-        self.rid_to_state = {}  # Dict[str -> ReqState]
+        self.rid_to_state: Dict[str, ReqState] = {}
 
     async def get_pixel_values(self, image_data):
         aspect_ratio = getattr(self.hf_config, "image_aspect_ratio", None)
@@ -153,10 +113,11 @@ class TokenizerManager:
                 image_data, aspect_ratio, grid_pinpoints, self.processor
             )
 
-    async def generate_request(self, obj: GenerateReqInput):
+    async def generate_request(self, obj: GenerateReqInput, request=None):
         if self.to_create_loop:
-            await self.create_handle_loop()
+            self.create_handle_loop()
 
+        obj.post_init()
         is_single = obj.is_single
         if is_single:
             rid = obj.rid
@@ -169,7 +130,7 @@ class TokenizerManager:
             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)"
+                    f"model's context length ({self.context_len} tokens)."
                 )
 
             sampling_params = SamplingParams(**obj.sampling_params)
@@ -207,23 +168,38 @@ class TokenizerManager:
             self.rid_to_state[rid] = state
 
             while True:
-                await event.wait()
-                out = self.convert_logprob_style(state.out_list[-1],
-                                                 obj.return_logprob,
-                                                 obj.top_logprobs_num,
-                                                 obj.return_text_in_logprobs)
+                try:
+                    await asyncio.wait_for(event.wait(), timeout=4)
+                except asyncio.TimeoutError:
+                    if request is not None and await request.is_disconnected():
+                        self.abort_request(rid)
+                        raise ValueError(f"Abort request {rid}")
+                    continue
+
+                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]
+
+                    yield out
+
                     break
+
                 event.clear()
+
+                yield out
         else:
-            assert obj.stream is False
+            if obj.stream:
+                raise ValueError("Do not support stream for batch mode.")
 
             if obj.input_ids is None:
                 bs = len(obj.text)
@@ -273,45 +249,83 @@ class TokenizerManager:
             for i in range(bs):
                 rid = obj.rid[i]
                 state = self.rid_to_state[rid]
-                await state.event.wait()
+
+                while True:
+                    try:
+                        await asyncio.wait_for(state.event.wait(), timeout=4)
+                        break
+                    except asyncio.TimeoutError:
+                        if request is not None and await request.is_disconnected():
+                            for rid in obj.rid:
+                                self.abort_request(rid)
+                            raise ValueError(f"Abort request {rid}")
+                        continue
+
                 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))
+                    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 flush_cache(self):
-        flush_cache_req = FlushCacheReq()
-        self.send_to_router.send_pyobj(flush_cache_req)
+    def flush_cache(self):
+        req = FlushCacheReq()
+        self.send_to_router.send_pyobj(req)
+
+    def abort_request(self, rid):
+        if rid not in self.rid_to_state:
+            return
+        del self.rid_to_state[rid]
+        req = AbortReq(rid)
+        self.send_to_router.send_pyobj(req)
+
+    def create_abort_task(self, obj: GenerateReqInput):
+        # Abort the request if the client is disconnected.
+        async def abort_request():
+            await asyncio.sleep(3)
+            if obj.is_single:
+                self.abort_request(obj.rid)
+            else:
+                for rid in obj.rids:
+                    self.abort_request(rid)
 
-    async def create_handle_loop(self):
+        background_tasks = BackgroundTasks()
+        background_tasks.add_task(abort_request)
+        return background_tasks
+
+    def create_handle_loop(self):
         self.to_create_loop = False
         loop = asyncio.get_event_loop()
         loop.create_task(self.handle_loop())
 
     async def handle_loop(self):
         while True:
-            recv_obj = await self.recv_from_detokenizer.recv_pyobj()
-
-            if isinstance(recv_obj, BatchStrOut):
-                for i, rid in enumerate(recv_obj.rids):
-                    recv_obj.meta_info[i]["id"] = rid
-                    out_dict = {
-                        "text": recv_obj.output_str[i],
-                        "meta_info": recv_obj.meta_info[i],
-                    }
-                    state = self.rid_to_state[rid]
-                    state.out_list.append(out_dict)
-                    state.finished = recv_obj.finished[i]
-                    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):
+            recv_obj: BatchTokenIDOut = await self.recv_from_detokenizer.recv_pyobj()
+            assert isinstance(recv_obj, BatchStrOut)
+
+            for i, rid in enumerate(recv_obj.rids):
+                state = self.rid_to_state.get(rid, None)
+                if state is None:
+                    continue
+
+                recv_obj.meta_info[i]["id"] = rid
+                out_dict = {
+                    "text": recv_obj.output_str[i],
+                    "meta_info": recv_obj.meta_info[i],
+                }
+                state.out_list.append(out_dict)
+                state.finished = recv_obj.finished_reason[i] is not None
+                state.event.set()
+
+    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
@@ -320,10 +334,14 @@ class TokenizerManager:
                 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"
+            ] = 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"
+            ] = self.detokenize_top_logprobs_tokens(
                 ret["meta_info"]["decode_top_logprobs"], return_text_in_logprobs
             )
         return ret
@@ -344,3 +362,49 @@ class TokenizerManager:
             if t:
                 top_logprobs[i] = self.detokenize_logprob_tokens(t, decode_to_text)
         return top_logprobs
+
+
+global global_processor
+
+
+def init_global_processor(server_args: ServerArgs):
+    global global_processor
+    transformers.logging.set_verbosity_error()
+    global_processor = get_processor(
+        server_args.tokenizer_path,
+        tokenizer_mode=server_args.tokenizer_mode,
+        trust_remote_code=server_args.trust_remote_code,
+    )
+
+
+def get_pixel_values(
+    image_data, image_aspect_ratio=None, image_grid_pinpoints=None, processor=None
+):
+    try:
+        processor = processor or global_processor
+        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:
+            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())

sglang/srt/model_config.py

@@ -1,5 +1,7 @@
 from typing import Optional
 
+from transformers import PretrainedConfig
+
 from sglang.srt.hf_transformers_utils import get_config, get_context_length
 
 
@@ -15,11 +17,14 @@ class ModelConfig:
         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)
-
+        self.model_overide_args = model_overide_args
+        self.hf_config = get_config(
+            self.path,
+            trust_remote_code,
+            revision,
+            model_overide_args=model_overide_args,
+        )
+        self.hf_text_config = get_hf_text_config(self.hf_config)
         if context_length is not None:
             self.context_len = context_length
         else:
@@ -45,3 +50,76 @@ class ModelConfig:
         self.hidden_size = self.hf_config.hidden_size
         self.num_hidden_layers = self.hf_config.num_hidden_layers
         self.vocab_size = self.hf_config.vocab_size
+
+    # adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py#L289
+    def get_total_num_kv_heads(self) -> int:
+        """Returns the total number of KV heads."""
+        # For GPTBigCode & Falcon:
+        # NOTE: for falcon, when new_decoder_architecture is True, the
+        # multi_query flag is ignored and we use n_head_kv for the number of
+        # KV heads.
+        falcon_model_types = ["falcon", "RefinedWeb", "RefinedWebModel"]
+        new_decoder_arch_falcon = (
+            self.hf_config.model_type in falcon_model_types
+            and getattr(self.hf_config, "new_decoder_architecture", False)
+        )
+        if not new_decoder_arch_falcon and getattr(
+            self.hf_text_config, "multi_query", False
+        ):
+            # Multi-query attention, only one KV head.
+            # Currently, tensor parallelism is not supported in this case.
+            return 1
+
+        # For DBRX and MPT
+        if self.hf_config.model_type in ["mpt"]:
+            if "kv_n_heads" in self.hf_config.attn_config:
+                return self.hf_config.attn_config["kv_n_heads"]
+            return self.hf_config.num_attention_heads
+        if self.hf_config.model_type in ["dbrx"]:
+            return getattr(
+                self.hf_config.attn_config,
+                "kv_n_heads",
+                self.hf_config.num_attention_heads,
+            )
+
+        attributes = [
+            # For Falcon:
+            "n_head_kv",
+            "num_kv_heads",
+            # For LLaMA-2:
+            "num_key_value_heads",
+            # For ChatGLM:
+            "multi_query_group_num",
+        ]
+        for attr in attributes:
+            num_kv_heads = getattr(self.hf_text_config, attr, None)
+            if num_kv_heads is not None:
+                return num_kv_heads
+
+        # For non-grouped-query attention models, the number of KV heads is
+        # equal to the number of attention heads.
+        return self.hf_text_config.num_attention_heads
+
+    # adapted from https://github.com/vllm-project/vllm/blob/main/vllm/config.py#L328
+    def get_num_kv_heads(self, tensor_parallel_size) -> int:
+        """Returns the number of KV heads per GPU."""
+        total_num_kv_heads = self.get_total_num_kv_heads()
+        # If tensor parallelism is used, we divide the number of KV heads by
+        # the tensor parallel size. We will replicate the KV heads in the
+        # case where the number of KV heads is smaller than the tensor
+        # parallel size so each GPU has at least one KV head.
+        return max(1, total_num_kv_heads // tensor_parallel_size)
+
+
+def get_hf_text_config(config: PretrainedConfig):
+    """Get the "sub" config relevant to llm for multi modal models.
+    No op for pure text models.
+    """
+    if hasattr(config, "text_config"):
+        # The code operates under the assumption that text_config should have
+        # `num_attention_heads` (among others). Assert here to fail early
+        # if transformers config doesn't align with this assumption.
+        assert hasattr(config.text_config, "num_attention_heads")
+        return config.text_config
+    else:
+        return config