sglang 0.1.15__py3-none-any.whl → 0.1.17__py3-none-any.whl

sglang/srt/managers/router/model_runner.py

@@ -1,30 +1,25 @@
 import importlib
 import importlib.resources
-import inspect
 import logging
 import pkgutil
 from dataclasses import dataclass
 from functools import lru_cache
-from typing import List
+from typing import List, Optional, Type
 
 import numpy as np
 import torch
-from vllm.model_executor.layers.quantization.awq import AWQConfig
-from vllm.model_executor.layers.quantization.gptq import GPTQConfig
-from vllm.model_executor.layers.quantization.marlin import MarlinConfig
-from vllm.model_executor.model_loader.utils import set_default_torch_dtype
+import torch.nn as nn
+from vllm.config import DeviceConfig, LoadConfig
+from vllm.config import ModelConfig as VllmModelConfig
 from vllm.distributed import initialize_model_parallel
+from vllm.model_executor.model_loader import get_model
+from vllm.model_executor.models import ModelRegistry
 
 from sglang.srt.managers.router.infer_batch import Batch, ForwardMode
 from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
-from sglang.srt.utils import is_multimodal_model
-from sglang.utils import get_available_gpu_memory
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import get_available_gpu_memory, is_multimodal_model
 
-QUANTIZATION_CONFIG_MAPPING = {
-    "awq": AWQConfig,
-    "gptq": GPTQConfig,
-    "marlin": MarlinConfig,
-}
 
 logger = logging.getLogger("model_runner")
 
@@ -32,35 +27,6 @@ logger = logging.getLogger("model_runner")
 global_server_args_dict = {}
 
 
-@lru_cache()
-def import_model_classes():
-    model_arch_name_to_cls = {}
-    package_name = "sglang.srt.models"
-    package = importlib.import_module(package_name)
-    for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."):
-        if not ispkg:
-            module = importlib.import_module(name)
-            if hasattr(module, "EntryClass"):
-                model_arch_name_to_cls[module.EntryClass.__name__] = module.EntryClass
-    return model_arch_name_to_cls
-
-
-def get_model_cls_by_arch_name(model_arch_names):
-    model_arch_name_to_cls = import_model_classes()
-
-    model_class = None
-    for arch in model_arch_names:
-        if arch in model_arch_name_to_cls:
-            model_class = model_arch_name_to_cls[arch]
-            break
-    else:
-        raise ValueError(
-            f"Unsupported architectures: {arch}. "
-            f"Supported list: {list(model_arch_name_to_cls.keys())}"
-        )
-    return model_class
-
-
 @dataclass
 class InputMetadata:
     model_runner: "ModelRunner"
@@ -110,8 +76,8 @@ class InputMetadata:
         self.kv_last_page_len = torch.ones(
             (self.batch_size,), dtype=torch.int32, device="cuda"
         )
-        req_pool_indices_cpu = self.req_pool_indices.cpu().tolist()
-        seq_lens_cpu = self.seq_lens.tolist()
+        req_pool_indices_cpu = self.req_pool_indices.cpu().numpy()
+        seq_lens_cpu = self.seq_lens.cpu().numpy()
         self.kv_indices = torch.cat(
             [
                 self.req_to_token_pool.req_to_token[
@@ -143,7 +109,7 @@ class InputMetadata:
                 self.kv_last_page_len,
                 self.model_runner.model_config.num_attention_heads // tp_size,
                 self.model_runner.model_config.num_key_value_heads // tp_size,
-                self.model_runner.model_config.head_dim
+                self.model_runner.model_config.head_dim,
             ]
 
             self.prefill_wrapper.begin_forward(*args)
@@ -253,113 +219,102 @@ class ModelRunner:
         tp_rank,
         tp_size,
         nccl_port,
-        load_format="auto",
-        trust_remote_code=True,
-        server_args_dict: dict = {},
+        server_args: ServerArgs,
     ):
         self.model_config = model_config
         self.mem_fraction_static = mem_fraction_static
         self.tp_rank = tp_rank
         self.tp_size = tp_size
         self.nccl_port = nccl_port
-        self.load_format = load_format
-        self.trust_remote_code = trust_remote_code
+        self.server_args = server_args
 
         global global_server_args_dict
-        global_server_args_dict = server_args_dict
+        global_server_args_dict = {
+            "enable_flashinfer": server_args.enable_flashinfer,
+            "attention_reduce_in_fp32": server_args.attention_reduce_in_fp32,
+        }
 
         # Init torch distributed
+        logger.info(f"[rank={self.tp_rank}] Set cuda device.")
         torch.cuda.set_device(self.tp_rank)
+        logger.info(f"[rank={self.tp_rank}] Init torch begin. Avail mem={get_available_gpu_memory(self.tp_rank):.2f} GB")
         torch.distributed.init_process_group(
             backend="nccl",
             world_size=self.tp_size,
             rank=self.tp_rank,
             init_method=f"tcp://127.0.0.1:{self.nccl_port}",
         )
-
         initialize_model_parallel(tensor_model_parallel_size=self.tp_size)
+        logger.info(f"[rank={self.tp_rank}] Init torch end.")
+
+        total_gpu_memory = get_available_gpu_memory(self.tp_rank, distributed=self.tp_size > 1)
+
+        if self.tp_size > 1:
+            total_local_gpu_memory = get_available_gpu_memory(self.tp_rank)
+            if total_local_gpu_memory < total_gpu_memory * 0.9:
+                raise ValueError("The memory capacity is unbalanced. Some GPUs may be occupied by other processes.")
 
-        total_gpu_memory = get_available_gpu_memory(
-            self.tp_rank, distributed=self.tp_size > 1
-        ) * (1 << 30)
         self.load_model()
         self.init_memory_pool(total_gpu_memory)
 
         self.is_multimodal_model = is_multimodal_model(self.model_config)
 
     def load_model(self):
-        """See also vllm/model_executor/model_loader.py::get_model"""
-        # Select model class
-        architectures = getattr(self.model_config.hf_config, "architectures", [])
-        model_class = get_model_cls_by_arch_name(architectures)
-        logger.info(f"Rank {self.tp_rank}: load weight begin.")
-
-        # Load weights
-        quant_config = None
-
-        quant_cfg = getattr(self.model_config.hf_config, "quantization_config", None)
-        if quant_cfg is not None:
-            quant_method = quant_cfg.get("quant_method", "").lower()
-            # compat: autogptq >=0.8.0 use checkpoint_format: str
-            # compat: autogptq <=0.7.1 is_marlin_format: bool
-            is_format_marlin = quant_cfg.get(
-                "checkpoint_format"
-            ) == "marlin" or quant_cfg.get("is_marlin_format", False)
-
-            # Use marlin if the GPTQ model is serialized in marlin format.
-            if quant_method == "gptq" and is_format_marlin:
-                quant_method = "marlin"
-
-            quant_config_class = QUANTIZATION_CONFIG_MAPPING.get(quant_method)
-
-            if quant_config_class is None:
-                raise ValueError(f"Unsupported quantization method: {quant_method}")
-
-            quant_config = quant_config_class.from_config(quant_cfg)
-            logger.info(f"quant_config: {quant_config}")
-
-        with set_default_torch_dtype(torch.float16):
-            with torch.device("cuda"):
-                model = model_class(
-                    config=self.model_config.hf_config, quant_config=quant_config
-                )
-            model.load_weights(
-                self.model_config.path,
-                cache_dir=None,
-                load_format=self.load_format,
-                revision=None,
-            )
-        self.model = model.eval()
-
-        logger.info(f"Rank {self.tp_rank}: load weight end.")
+        logger.info(f"[rank={self.tp_rank}] Load weight begin.")
+
+        device_config = DeviceConfig()
+        load_config = LoadConfig(load_format=self.server_args.load_format)
+        vllm_model_config = VllmModelConfig(
+            model=self.server_args.model_path,
+            quantization=self.server_args.quantization,
+            tokenizer=None,
+            tokenizer_mode=None,
+            trust_remote_code=self.server_args.trust_remote_code,
+            dtype=torch.float16,
+            seed=42,
+            skip_tokenizer_init=True,
+        )
+        if self.model_config.model_overide_args is not None:
+            vllm_model_config.hf_config.update(self.model_config.model_overide_args)
+
+        self.model = get_model(
+            model_config=vllm_model_config,
+            device_config=device_config,
+            load_config=load_config,
+            lora_config=None,
+            vision_language_config=None,
+            parallel_config=None,
+            scheduler_config=None,
+        )
+        logger.info(f"[rank={self.tp_rank}] Load weight end. "
+                    f"Type={type(self.model).__name__}. "
+                    f"Avail mem={get_available_gpu_memory(self.tp_rank):.2f} GB")
 
     def profile_max_num_token(self, total_gpu_memory):
-        available_gpu_memory = get_available_gpu_memory(
-            self.tp_rank, distributed=self.tp_size > 1
-        ) * (1 << 30)
+        available_gpu_memory = get_available_gpu_memory(self.tp_rank, distributed=self.tp_size > 1)
         head_dim = self.model_config.head_dim
         head_num = self.model_config.num_key_value_heads // self.tp_size
         cell_size = head_num * head_dim * self.model_config.num_hidden_layers * 2 * 2
         rest_memory = available_gpu_memory - total_gpu_memory * (
             1 - self.mem_fraction_static
         )
-        max_num_token = int(rest_memory // cell_size)
+        max_num_token = int(rest_memory * (1 << 30) // cell_size)
         return max_num_token
 
     def init_memory_pool(self, total_gpu_memory):
-        self.max_total_num_token = self.profile_max_num_token(total_gpu_memory)
+        self.max_total_num_tokens = self.profile_max_num_token(total_gpu_memory)
 
-        if self.max_total_num_token <= 0:
+        if self.max_total_num_tokens <= 0:
             raise RuntimeError(
                 "Not enought memory. " "Please try to increase --mem-fraction-static."
             )
 
         self.req_to_token_pool = ReqToTokenPool(
-            int(self.max_total_num_token / self.model_config.context_len * 256),
+            int(self.max_total_num_tokens / self.model_config.context_len * 256),
             self.model_config.context_len + 8,
         )
         self.token_to_kv_pool = TokenToKVPool(
-            self.max_total_num_token,
+            self.max_total_num_tokens,
             dtype=torch.float16,
             head_num=self.model_config.num_key_value_heads // self.tp_size,
             head_dim=self.model_config.head_dim,
@@ -456,3 +411,35 @@ class ModelRunner:
             return self.forward_prefill(batch)
         else:
             raise ValueError(f"Invaid forward mode: {forward_mode}")
+
+
+@lru_cache()
+def import_model_classes():
+    model_arch_name_to_cls = {}
+    package_name = "sglang.srt.models"
+    package = importlib.import_module(package_name)
+    for _, name, ispkg in pkgutil.iter_modules(package.__path__, package_name + "."):
+        if not ispkg:
+            module = importlib.import_module(name)
+            if hasattr(module, "EntryClass"):
+                entry = module.EntryClass
+                if isinstance(entry, list): # To support multiple model classes in one module
+                    for cls in entry:
+                        model_arch_name_to_cls[cls.__name__] = cls
+                else:
+                    model_arch_name_to_cls[entry.__name__] = entry
+    return model_arch_name_to_cls
+
+
+def load_model_cls_srt(model_arch: str) -> Optional[Type[nn.Module]]:
+    model_arch_name_to_cls = import_model_classes()
+    if model_arch not in model_arch_name_to_cls:
+        raise ValueError(
+            f"Unsupported architectures: {model_arch}. "
+            f"Supported list: {list(model_arch_name_to_cls.keys())}"
+        )
+    return model_arch_name_to_cls[model_arch]
+
+
+# Monkey patch model loader
+setattr(ModelRegistry, "load_model_cls", load_model_cls_srt)

sglang/srt/managers/router/radix_cache.py

@@ -11,7 +11,7 @@ class TreeNode:
         self.parent = None
         self.key = None
         self.value = None
-        self.ref_counter = 0
+        self.lock_ref = 0
         self.last_access_time = time.time()
 
     def __lt__(self, other: "TreeNode"):
@@ -28,7 +28,9 @@ def _key_match(key0, key1):
 
 
 class RadixCache:
-    def __init__(self, disable: bool = False):
+    def __init__(self, req_to_token_pool, token_to_kv_pool, disable: bool = False):
+        self.req_to_token_pool = req_to_token_pool
+        self.token_to_kv_pool = token_to_kv_pool
         self.disable = disable
         self.reset()
 
@@ -38,7 +40,7 @@ class RadixCache:
         self.root_node = TreeNode()
         self.root_node.key = []
         self.root_node.value = []
-        self.root_node.ref_counter = 1
+        self.root_node.lock_ref = 1
         self.evictable_size_ = 0
 
     def match_prefix(self, key):
@@ -50,16 +52,52 @@ class RadixCache:
         self._match_prefix_helper(self.root_node, key, value, last_node)
         if value:
             value = torch.concat(value)
+        else:
+            value = torch.tensor([], dtype=torch.int64)
         return value, last_node[0]
 
     def insert(self, key, value=None):
         if self.disable:
-            return len(key)
+            return 0
 
         if value is None:
             value = [x for x in key]
         return self._insert_helper(self.root_node, key, value)
 
+    def cache_req(
+        self,
+        token_ids,
+        last_uncached_pos,
+        req_pool_idx,
+        del_in_memory_pool=True,
+        old_last_node=None,
+    ):
+        # Insert the request into radix cache
+        indices = self.req_to_token_pool.req_to_token[req_pool_idx, : len(token_ids)]
+        new_prefix_len = self.insert(token_ids, indices.clone())
+
+        if self.disable:
+            if del_in_memory_pool:
+                self.token_to_kv_pool.dec_refs(indices)
+            else:
+                return torch.tensor([], dtype=torch.int64), self.root_node
+
+        # Radix Cache takes one ref in memory pool
+        self.token_to_kv_pool.dec_refs(indices[last_uncached_pos:new_prefix_len])
+
+        if del_in_memory_pool:
+            self.req_to_token_pool.free(req_pool_idx)
+        else:
+            cached_indices, new_last_node = self.match_prefix(token_ids)
+            assert len(cached_indices) == len(token_ids)
+
+            self.req_to_token_pool.req_to_token[
+                req_pool_idx, last_uncached_pos : len(cached_indices)
+            ] = cached_indices[last_uncached_pos:]
+            self.dec_lock_ref(old_last_node)
+            self.inc_lock_ref(new_last_node)
+            return cached_indices, new_last_node
+
     def pretty_print(self):
         self._print_helper(self.root_node, 0)
         print(f"#tokens: {self.total_size()}")
@@ -80,7 +118,7 @@ class RadixCache:
 
             if x == self.root_node:
                 break
-            if x.ref_counter > 0:
+            if x.lock_ref > 0:
                 continue
 
             num_evicted += evict_callback(x.value)
@@ -89,23 +127,23 @@ class RadixCache:
             if len(x.parent.children) == 0:
                 heapq.heappush(leaves, x.parent)
 
-    def inc_ref_counter(self, node):
+    def inc_lock_ref(self, node: TreeNode):
         delta = 0
         while node != self.root_node:
-            if node.ref_counter == 0:
+            if node.lock_ref == 0:
                 self.evictable_size_ -= len(node.value)
                 delta -= len(node.value)
-            node.ref_counter += 1
+            node.lock_ref += 1
             node = node.parent
         return delta
 
-    def dec_ref_counter(self, node):
+    def dec_lock_ref(self, node: TreeNode):
         delta = 0
         while node != self.root_node:
-            if node.ref_counter == 1:
+            if node.lock_ref == 1:
                 self.evictable_size_ += len(node.value)
                 delta += len(node.value)
-            node.ref_counter -= 1
+            node.lock_ref -= 1
             node = node.parent
         return delta
 
@@ -131,12 +169,12 @@ class RadixCache:
                 last_node[0] = child
                 self._match_prefix_helper(child, key[prefix_len:], value, last_node)
 
-    def _split_node(self, key, child, split_len):
+    def _split_node(self, key, child: TreeNode, split_len):
         # new_node -> child
         new_node = TreeNode()
         new_node.children = {key[split_len:][0]: child}
         new_node.parent = child.parent
-        new_node.ref_counter = child.ref_counter
+        new_node.lock_ref = child.lock_ref
         new_node.key = child.key[:split_len]
         new_node.value = child.value[:split_len]
         child.parent = new_node
@@ -176,11 +214,9 @@ class RadixCache:
             self.evictable_size_ += len(value)
         return 0
 
-    def _print_helper(self, node, indent):
+    def _print_helper(self, node: TreeNode, indent):
         for _, child in node.children.items():
-            print(
-                " " * indent, len(child.key), child.key[:10], f"r={child.ref_counter}"
-            )
+            print(" " * indent, len(child.key), child.key[:10], f"r={child.lock_ref}")
             self._print_helper(child, indent=indent + 2)
 
     def _delete_leaf(self, node):
@@ -211,7 +247,7 @@ class RadixCache:
 
 
 if __name__ == "__main__":
-    tree = RadixCache()
+    tree = RadixCache(None, None, False)
 
     tree.insert("Hello")
     tree.insert("Hello")

sglang/srt/managers/router/scheduler.py

@@ -6,15 +6,15 @@ class Scheduler:
     def __init__(
         self,
         schedule_heuristic,
-        max_running_seq,
-        max_prefill_num_token,
-        max_total_num_token,
+        max_running_seqs,
+        max_prefill_num_tokens,
+        max_total_num_tokens,
         tree_cache,
     ):
         self.schedule_heuristic = schedule_heuristic
-        self.max_running_seq = max_running_seq
-        self.max_prefill_num_token = max_prefill_num_token
-        self.max_total_num_token = max_total_num_token
+        self.max_running_seqs = max_running_seqs
+        self.max_prefill_num_tokens = max_prefill_num_tokens
+        self.max_total_num_tokens = max_total_num_tokens
         self.tree_cache = tree_cache
 
     def get_priority_queue(self, forward_queue):
@@ -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])