sglang 0.3.1.post3__py3-none-any.whl → 0.3.3__py3-none-any.whl

sglang/srt/lora/lora_manager.py

@@ -17,21 +17,56 @@ limitations under the License.
 # and "Punica: Multi-Tenant LoRA Serving"
 
 
+import logging
 import re
-from dataclasses import dataclass
 
 import torch
 
 from sglang.srt.lora.lora import LoRAAdapter, get_lora_layer
 from sglang.srt.lora.lora_config import LoRAConfig
-from sglang.srt.model_executor.forward_batch_info import ForwardMode
-from sglang.srt.utils import is_hip, replace_submodule
+from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+from sglang.srt.utils import is_flashinfer_available, replace_submodule
 
-# ROCm: flashinfer available later
-if not is_hip():
+logger = logging.getLogger(__name__)
+
+if is_flashinfer_available():
     from flashinfer import SegmentGEMMWrapper
 
 
+def get_module_name(name):
+    # Fallback solution of mapping from config module name to module name in model class.
+    # Please check if it aligns with your base model.
+    # Please implement the function in the model class if it is not.
+    # You can reference this function in llama.py.
+    params_mapping = {
+        "q_proj": "qkv_proj",
+        "k_proj": "qkv_proj",
+        "v_proj": "qkv_proj",
+        "gate_proj": "gate_up_proj",
+        "up_proj": "gate_up_proj",
+    }
+    return params_mapping.get(name, name)
+
+
+def get_hidden_dim(module_name, config):
+    # Fallback solution of get_hidden_dim for different modules
+    # Please check if it aligns with your base model.
+    # Please implement the function in the model class if it is not.
+    # You can reference this function in llama.py.
+    if module_name in ["q_proj", "o_proj", "qkv_proj"]:
+        return config.hidden_size, config.hidden_size
+    elif module_name in ["kv_proj"]:
+        return config.hidden_size, config.hidden_size // (
+            config.num_attention_heads // config.num_key_value_heads
+        )
+    elif module_name == "gate_up_proj":
+        return config.hidden_size, config.intermediate_size
+    elif module_name == "down_proj":
+        return config.intermediate_size, config.hidden_size
+    else:
+        raise NotImplementedError()
+
+
 def get_stacked_name(name):
     # origin name -> (name for A, name for B)
     params_mapping = {
@@ -104,12 +139,20 @@ class LoRAManager:
             self.origin_target_modules = set(self.origin_target_modules) | set(
                 self.configs[name].target_modules
             )
-        self.target_modules = set(
-            [
+        if hasattr(self.base_model, "get_module_name"):
+            self.target_modules = {
                 self.base_model.get_module_name(module)
                 for module in self.origin_target_modules
-            ]
-        )
+            }
+        else:
+            logger.warning(
+                f"WARNING: get_module_name() is not defined, "
+                f"which is used to map config module name to model implementation module name."
+                f"Use the default one, but please check if it is correct for your model."
+            )
+            self.target_modules = {
+                get_module_name(module) for module in self.origin_target_modules
+            }
         self.target_weights = set(
             [get_stacked_name(module) for module in self.origin_target_modules]
         )
@@ -147,7 +190,15 @@ class LoRAManager:
         num_layer = self.base_hf_config.num_hidden_layers
         for module_A, module_B in self.target_weights:
             # init A tensor, column_major=True
-            hidden_dim_A, _ = self.base_model.get_hidden_dim(module_A)
+            if hasattr(self.base_model, "get_hidden_dim"):
+                hidden_dim_A, _ = self.base_model.get_hidden_dim(module_A)
+            else:
+                logger.warning(
+                    f"WARNING: get_hidden_dim() is not defined, "
+                    f"which is used to get the hidden dim for different lora modules"
+                    f"Use the default one, but please check if it is correct for your model."
+                )
+                hidden_dim_A, _ = get_hidden_dim(module_A, self.base_hf_config)
             c = self.loras[-1].get_stacked_multiply(module_A)
             if module_A not in self.A_buffer:
                 self.A_buffer[module_A] = [
@@ -163,7 +214,15 @@ class LoRAManager:
                     for i in range(num_layer)
                 ]
             # init B tensor, column_major=True
-            _, hidden_dim_B = self.base_model.get_hidden_dim(module_B)
+            if hasattr(self.base_model, "get_hidden_dim"):
+                _, hidden_dim_B = self.base_model.get_hidden_dim(module_B)
+            else:
+                logger.warning(
+                    f"WARNING: get_hidden_dim() is not defined, "
+                    f"which is used to get the hidden dim for different lora modules"
+                    f"Use the default one, but please check if it is correct for your model."
+                )
+                _, hidden_dim_B = get_hidden_dim(module_B, self.base_hf_config)
             c = self.loras[-1].get_stacked_multiply(module_B)
             if module_B not in self.B_buffer:
                 self.B_buffer[module_B] = [
@@ -208,33 +267,46 @@ class LoRAManager:
                     if lora_weight_name:
                         self.B_buffer[lora_weight_name][i][buffer_id].copy_(weights)
 
-    def prepare_lora_batch(self, batch, extend_seq_lens=None):
+    def prepare_lora_batch(self, forward_batch: ForwardBatch):
         # load active loras into lora memory pool
-        cur_uids = set([req.lora_path for req in batch.reqs])
+        cur_uids = set(forward_batch.lora_paths)
         assert len(cur_uids) <= self.max_loras_per_batch
         i = 0
+        j = len(self.active_uids)
         evictable_uids = list(self.active_uids)
         for uid in cur_uids:
             if uid not in self.active_uids:
-                while i < len(evictable_uids) and evictable_uids[i] in cur_uids:
-                    i += 1
-                if i < len(evictable_uids):
+                if j < self.max_loras_per_batch:
+                    index = j
+                    j += 1
+                else:
+                    while i < len(evictable_uids) and evictable_uids[i] in cur_uids:
+                        i += 1
+                    assert i < len(evictable_uids)
                     self.active_uids.remove(evictable_uids[i])
                     self.buffer_id.pop(evictable_uids[i])
-                self.load_lora(uid, i)
+                    index = i
+                    i += 1
+                self.load_lora(uid, index)
                 self.active_uids.add(uid)
-                self.buffer_id[uid] = i
-                i += 1
+                self.buffer_id[uid] = index
 
         if cur_uids == set([None]):
             return
 
         # setup lora in forward modules
-        bs = len(batch.reqs)
-        seg_lens = extend_seq_lens if batch.forward_mode.is_extend() else torch.ones(bs)
+        bs = forward_batch.batch_size
+        seg_lens = (
+            forward_batch.extend_seq_lens
+            if forward_batch.forward_mode.is_extend()
+            else torch.ones(bs, device="cuda")
+        )
+        # FIXME: reuse the data rather than recompute
+        seg_indptr = torch.zeros((bs + 1,), dtype=torch.int32, device="cuda")
+        seg_indptr[1:] = torch.cumsum(seg_lens, dim=0)
         weight_indices = torch.empty((bs,), dtype=torch.int64, device="cuda")
-        for i, req in enumerate(batch.reqs):
-            weight_indices[i] = self.buffer_id[req.lora_path]
+        for i, lora_path in enumerate(forward_batch.lora_paths):
+            weight_indices[i] = self.buffer_id[lora_path]
 
         for module_name, module in self.lora_modules:
             layer_id = get_layer_id(module_name)
@@ -245,7 +317,7 @@ class LoRAManager:
                     self.A_buffer[weight_name][layer_id],
                     self.B_buffer[weight_name][layer_id],
                     bs,
-                    seg_lens,
+                    seg_indptr,
                     weight_indices,
                 )
             else:
@@ -254,6 +326,6 @@ class LoRAManager:
                     self.B_buffer["q_proj"][layer_id],
                     self.B_buffer["kv_proj"][layer_id],
                     bs,
-                    seg_lens,
+                    seg_indptr,
                     weight_indices,
                 )

sglang/srt/managers/detokenizer_manager.py

@@ -15,13 +15,12 @@ limitations under the License.
 
 """DetokenizerManager is a process that detokenizes the token ids."""
 
-import asyncio
 import dataclasses
+import logging
+from collections import OrderedDict
 from typing import List
 
-import uvloop
 import zmq
-import zmq.asyncio
 
 from sglang.srt.hf_transformers_utils import get_tokenizer
 from sglang.srt.managers.io_struct import (
@@ -32,9 +31,10 @@ from sglang.srt.managers.io_struct import (
 )
 from sglang.srt.managers.schedule_batch import FINISH_MATCHED_STR
 from sglang.srt.server_args import PortArgs, ServerArgs
+from sglang.srt.utils import configure_logger, kill_parent_process
 from sglang.utils import find_printable_text, get_exception_traceback
 
-asyncio.set_event_loop_policy(uvloop.EventLoopPolicy())
+logger = logging.getLogger(__name__)
 
 
 @dataclasses.dataclass
@@ -57,12 +57,12 @@ class DetokenizerManager:
         port_args: PortArgs,
     ):
         # Init inter-process communication
-        context = zmq.asyncio.Context(2)
-        self.recv_from_router = context.socket(zmq.PULL)
-        self.recv_from_router.bind(f"tcp://127.0.0.1:{port_args.detokenizer_port}")
+        context = zmq.Context(2)
+        self.recv_from_scheduler = context.socket(zmq.PULL)
+        self.recv_from_scheduler.bind(f"ipc://{port_args.detokenizer_ipc_name}")
 
         self.send_to_tokenizer = context.socket(zmq.PUSH)
-        self.send_to_tokenizer.connect(f"tcp://127.0.0.1:{port_args.tokenizer_port}")
+        self.send_to_tokenizer.connect(f"ipc://{port_args.tokenizer_ipc_name}")
 
         if server_args.skip_tokenizer_init:
             self.tokenizer = None
@@ -73,13 +73,13 @@ class DetokenizerManager:
                 trust_remote_code=server_args.trust_remote_code,
             )
 
-        self.decode_status = {}
+        self.decode_status = LimitedCapacityDict()
 
-    async def handle_loop(self):
+    def event_loop(self):
         """The event loop that handles requests"""
 
         while True:
-            recv_obj = await self.recv_from_router.recv_pyobj()
+            recv_obj = self.recv_from_scheduler.recv_pyobj()
 
             if isinstance(recv_obj, BatchEmbeddingOut):
                 # If it is embedding model, no detokenization is needed.
@@ -170,16 +170,29 @@ class DetokenizerManager:
             )
 
 
-def start_detokenizer_process(
+class LimitedCapacityDict(OrderedDict):
+    def __init__(self, capacity=1 << 15, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.capacity = capacity
+
+    def __setitem__(self, key, value):
+        if len(self) >= self.capacity:
+            # Remove the oldest element (first item in the dict)
+            self.popitem(last=False)
+        # Set the new item
+        super().__setitem__(key, value)
+
+
+def run_detokenizer_process(
     server_args: ServerArgs,
     port_args: PortArgs,
-    pipe_writer,
 ):
+    configure_logger(server_args)
+
     try:
         manager = DetokenizerManager(server_args, port_args)
+        manager.event_loop()
     except Exception:
-        pipe_writer.send(get_exception_traceback())
-        raise
-    pipe_writer.send("init ok")
-    loop = asyncio.get_event_loop()
-    loop.run_until_complete(manager.handle_loop())
+        msg = get_exception_traceback()
+        logger.error(msg)
+        kill_parent_process()

sglang/srt/managers/image_processor.py

@@ -0,0 +1,187 @@
+# TODO: also move pad_input_ids into this module
+import asyncio
+import concurrent.futures
+import logging
+import multiprocessing as mp
+import os
+from abc import ABC, abstractmethod
+from typing import List, Optional, Union
+
+import numpy as np
+import transformers
+
+from sglang.srt.hf_transformers_utils import get_processor
+from sglang.srt.mm_utils import expand2square, process_anyres_image
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import load_image
+from sglang.utils import get_exception_traceback
+
+logger = logging.getLogger(__name__)
+
+global global_processor
+
+
+def init_global_processor(server_args: ServerArgs):
+    """Init the global processor for multi modal models."""
+    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,
+    )
+
+
+class BaseImageProcessor(ABC):
+    @abstractmethod
+    async def process_images_async(self, image_data, **kwargs):
+        pass
+
+
+class DummyImageProcessor(BaseImageProcessor):
+    async def process_images_async(self, *args, **kwargs):
+        return None
+
+
+class LlavaImageProcessor(BaseImageProcessor):
+    def __init__(self, hf_config, server_args, _image_processor):
+        self.hf_config = hf_config
+        self._image_processor = _image_processor
+        self.executor = concurrent.futures.ProcessPoolExecutor(
+            initializer=init_global_processor,
+            mp_context=mp.get_context("fork"),
+            initargs=(server_args,),
+            max_workers=os.environ.get("SGLANG_CPU_COUNT", os.cpu_count()),
+        )
+
+    @staticmethod
+    def _process_single_image_task(
+        image_data: Union[str, bytes],
+        image_aspect_ratio: Optional[str] = None,
+        image_grid_pinpoints: Optional[str] = None,
+        image_processor=None,
+    ):
+        image_processor = image_processor or global_processor.image_processor
+
+        try:
+            image, image_size = load_image(image_data)
+            if image_size is not None:
+                # It is a video with multiple images
+                image_hash = hash(image_data)
+                pixel_values = 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:
+                # It is an image
+                image_hash = hash(image_data)
+                if image_aspect_ratio == "pad":
+                    image = expand2square(
+                        image,
+                        tuple(int(x * 255) for x in image_processor.image_mean),
+                    )
+                    pixel_values = image_processor(image.convert("RGB"))[
+                        "pixel_values"
+                    ][0]
+                elif image_aspect_ratio == "anyres" or (
+                    image_aspect_ratio is not None
+                    and "anyres_max" in image_aspect_ratio
+                ):
+                    pixel_values = process_anyres_image(
+                        image, image_processor, image_grid_pinpoints
+                    )
+                else:
+                    pixel_values = image_processor(image)["pixel_values"][0]
+
+                if isinstance(pixel_values, np.ndarray):
+                    pixel_values = pixel_values.astype(np.float16)
+
+                return pixel_values, image_hash, image.size
+        except Exception:
+            logger.error("Exception in TokenizerManager:\n" + get_exception_traceback())
+
+    async def _process_single_image(
+        self, image_data: Union[bytes, str], aspect_ratio: str, grid_pinpoints: str
+    ):
+        if self.executor is not None:
+            loop = asyncio.get_event_loop()
+            return await loop.run_in_executor(
+                self.executor,
+                LlavaImageProcessor._process_single_image_task,
+                image_data,
+                aspect_ratio,
+                grid_pinpoints,
+            )
+        else:
+            return self._process_single_image_task(
+                image_data, aspect_ratio, grid_pinpoints
+            )
+
+    async def process_images_async(
+        self, image_data: List[Union[str, bytes]], request_obj
+    ):
+        if not image_data:
+            return None
+
+        aspect_ratio = getattr(self.hf_config, "image_aspect_ratio", None)
+        grid_pinpoints = (
+            self.hf_config.image_grid_pinpoints
+            if hasattr(self.hf_config, "image_grid_pinpoints")
+            and "anyres" in aspect_ratio
+            else None
+        )
+
+        if isinstance(image_data, list) and len(image_data) > 0:
+            # Multiple images
+            if len(image_data) > 1:
+                aspect_ratio = "pad"  # LLaVA OneVision Handling: more than one image --> interleaved image mode or video mode. We do not use anyres
+                pixel_values, image_hashes, image_sizes = [], [], []
+                res = []
+                for img_data in image_data:
+                    res.append(
+                        self._process_single_image(
+                            img_data, aspect_ratio, grid_pinpoints
+                        )
+                    )
+                res = await asyncio.gather(*res)
+                for pixel_v, image_h, image_s in res:
+                    pixel_values.append(pixel_v)
+                    image_hashes.append(image_h)
+                    image_sizes.append(image_s)
+
+                if isinstance(pixel_values[0], np.ndarray):
+                    pixel_values = np.stack(pixel_values, axis=0)
+            else:
+                # A single image
+                pixel_values, image_hash, image_size = await self._process_single_image(
+                    image_data[0], aspect_ratio, grid_pinpoints
+                )
+                image_hashes = [image_hash]
+                image_sizes = [image_size]
+        elif isinstance(image_data, str):
+            # A single image
+            pixel_values, image_hash, image_size = await self._process_single_image(
+                image_data, aspect_ratio, grid_pinpoints
+            )
+            image_hashes = [image_hash]
+            image_sizes = [image_size]
+        else:
+            raise ValueError(f"Invalid image data: {image_data}")
+
+        return {
+            "pixel_values": pixel_values,
+            "image_hashes": image_hashes,
+            "image_sizes": image_sizes,
+            "modalities": request_obj.modalities,
+        }
+
+
+def get_image_processor(
+    hf_config, server_args: ServerArgs, _image_processor
+) -> BaseImageProcessor:
+    return LlavaImageProcessor(hf_config, server_args, _image_processor)
+
+
+def get_dummy_image_processor():
+    return DummyImageProcessor()