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

sglang/srt/managers/controller/model_runner.py

@@ -0,0 +1,359 @@
+"""ModelRunner runs the forward passes of the models."""
+
+import importlib
+import importlib.resources
+import logging
+import pkgutil
+from functools import lru_cache
+from typing import Optional, Type
+
+import torch
+import torch.nn as nn
+from vllm.config import DeviceConfig, LoadConfig
+from vllm.config import ModelConfig as VllmModelConfig
+from vllm.distributed import init_distributed_environment, initialize_model_parallel, get_tp_group
+from vllm.model_executor.model_loader import get_model
+from vllm.model_executor.models import ModelRegistry
+
+from sglang.global_config import global_config
+from sglang.srt.managers.controller.infer_batch import (
+    Batch,
+    ForwardMode,
+    InputMetadata,
+    global_server_args_dict,
+)
+from sglang.srt.memory_pool import ReqToTokenPool, TokenToKVPool
+from sglang.srt.server_args import ServerArgs
+from sglang.srt.utils import (
+    get_available_gpu_memory,
+    is_multimodal_model,
+    monkey_patch_vllm_dummy_weight_loader,
+    monkey_patch_vllm_p2p_access_check,
+)
+
+logger = logging.getLogger("srt.model_runner")
+
+
+class ModelRunner:
+    def __init__(
+        self,
+        model_config,
+        mem_fraction_static: float,
+        gpu_id: int,
+        tp_rank: int,
+        tp_size: int,
+        nccl_port: int,
+        server_args: ServerArgs,
+    ):
+        # Parse args
+        self.model_config = model_config
+        self.mem_fraction_static = mem_fraction_static
+        self.gpu_id = gpu_id
+        self.tp_rank = tp_rank
+        self.tp_size = tp_size
+        self.nccl_port = nccl_port
+        self.server_args = server_args
+        self.is_multimodal_model = is_multimodal_model(self.model_config)
+        monkey_patch_vllm_dummy_weight_loader()
+
+        # Init torch distributed
+        torch.cuda.set_device(self.gpu_id)
+        logger.info(f"[gpu_id={self.gpu_id}] Init nccl begin.")
+
+        if not server_args.enable_p2p_check:
+            monkey_patch_vllm_p2p_access_check(self.gpu_id)
+
+        if server_args.nccl_init_addr:
+            nccl_init_method = f"tcp://{server_args.nccl_init_addr}"
+        else:
+            nccl_init_method = f"tcp://127.0.0.1:{self.nccl_port}"
+        init_distributed_environment(
+            backend="nccl",
+            world_size=self.tp_size,
+            rank=self.tp_rank,
+            local_rank=self.gpu_id,
+            distributed_init_method=nccl_init_method,
+        )
+        initialize_model_parallel(tensor_model_parallel_size=self.tp_size)
+        self.tp_group = get_tp_group()
+        total_gpu_memory = get_available_gpu_memory(
+            self.gpu_id, distributed=self.tp_size > 1
+        )
+
+        if self.tp_size > 1:
+            total_local_gpu_memory = get_available_gpu_memory(self.gpu_id)
+            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."
+                )
+
+        # Set some global args
+        global_server_args_dict["disable_flashinfer"] = server_args.disable_flashinfer
+        global_server_args_dict[
+            "attention_reduce_in_fp32"
+        ] = server_args.attention_reduce_in_fp32
+
+        # Load the model and create memory pool
+        self.load_model()
+        self.init_memory_pool(total_gpu_memory)
+        self.init_cublas()
+        self.init_flash_infer()
+
+        # Capture cuda graphs
+        self.init_cuda_graphs()
+
+    def load_model(self):
+        logger.info(
+            f"[gpu_id={self.gpu_id}] Load weight begin. "
+            f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
+        )
+
+        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=self.server_args.dtype,
+            seed=42,
+            skip_tokenizer_init=True,
+        )
+        self.dtype = vllm_model_config.dtype
+        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,
+            multimodal_config=None,
+            parallel_config=None,
+            scheduler_config=None,
+            cache_config=None,
+        )
+        logger.info(
+            f"[gpu_id={self.gpu_id}] Load weight end. "
+            f"type={type(self.model).__name__}, "
+            f"dtype={self.dtype}, "
+            f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
+        )
+
+    def profile_max_num_token(self, total_gpu_memory):
+        available_gpu_memory = get_available_gpu_memory(
+            self.gpu_id, distributed=self.tp_size > 1
+        )
+        head_dim = self.model_config.head_dim
+        head_num = self.model_config.get_num_kv_heads(self.tp_size)
+        cell_size = (
+            head_num
+            * head_dim
+            * self.model_config.num_hidden_layers
+            * 2
+            * torch._utils._element_size(self.dtype)
+        )
+        rest_memory = available_gpu_memory - total_gpu_memory * (
+            1 - self.mem_fraction_static
+        )
+        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_tokens = self.profile_max_num_token(total_gpu_memory)
+
+        if self.max_total_num_tokens <= 0:
+            raise RuntimeError(
+                "Not enough memory. Please try to increase --mem-fraction-static."
+            )
+
+        self.req_to_token_pool = ReqToTokenPool(
+            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_tokens,
+            dtype=self.dtype,
+            head_num=self.model_config.get_num_kv_heads(self.tp_size),
+            head_dim=self.model_config.head_dim,
+            layer_num=self.model_config.num_hidden_layers,
+        )
+        logger.info(
+            f"[gpu_id={self.gpu_id}] Memory pool end. "
+            f"avail mem={get_available_gpu_memory(self.gpu_id):.2f} GB"
+        )
+
+    def init_cublas(self):
+        """We need to run a small matmul to init cublas. Otherwise, it will raise some errors later."""
+        dtype = torch.float16
+        device = "cuda"
+        a = torch.ones((16, 16), dtype=dtype, device=device)
+        b = torch.ones((16, 16), dtype=dtype, device=device)
+        c = a @ b
+        return c
+
+    def init_flash_infer(self):
+        if self.server_args.disable_flashinfer:
+            self.flashinfer_prefill_wrapper_ragged = None
+            self.flashinfer_prefill_wrapper_paged = None
+            self.flashinfer_decode_wrapper = None
+            return
+
+        from flashinfer import (
+            BatchDecodeWithPagedKVCacheWrapper,
+            BatchPrefillWithPagedKVCacheWrapper,
+            BatchPrefillWithRaggedKVCacheWrapper,
+        )
+        from flashinfer.decode import _grouped_size_compiled_for_decode_kernels
+
+        if not _grouped_size_compiled_for_decode_kernels(
+            self.model_config.num_attention_heads // self.tp_size,
+            self.model_config.get_num_kv_heads(self.tp_size),
+        ):
+            use_tensor_cores = True
+        else:
+            use_tensor_cores = False
+
+        self.flashinfer_workspace_buffers = torch.empty(
+            2, global_config.flashinfer_workspace_size, dtype=torch.uint8, device="cuda"
+        )
+        self.flashinfer_prefill_wrapper_ragged = BatchPrefillWithRaggedKVCacheWrapper(
+            self.flashinfer_workspace_buffers[0], "NHD"
+        )
+        self.flashinfer_prefill_wrapper_paged = BatchPrefillWithPagedKVCacheWrapper(
+            self.flashinfer_workspace_buffers[1], "NHD"
+        )
+        self.flashinfer_decode_wrapper = BatchDecodeWithPagedKVCacheWrapper(
+            self.flashinfer_workspace_buffers[0],
+            "NHD",
+            use_tensor_cores=use_tensor_cores,
+        )
+
+    def init_cuda_graphs(self):
+        from sglang.srt.managers.controller.cuda_graph_runner import CudaGraphRunner
+
+        if self.server_args.disable_cuda_graph or self.server_args.disable_flashinfer:
+            self.cuda_graph_runner = None
+            return
+
+        logger.info(f"[gpu_id={self.gpu_id}] Capture cuda graph begin.")
+        batch_size_list = [1, 2, 4] + [i * 8 for i in range(1, 16)]
+        self.cuda_graph_runner = CudaGraphRunner(
+            self, max_batch_size_to_capture=max(batch_size_list)
+        )
+        self.cuda_graph_runner.capture(batch_size_list)
+
+    @torch.inference_mode()
+    def forward_decode(self, batch: Batch):
+        if self.cuda_graph_runner and self.cuda_graph_runner.can_run(len(batch.reqs)):
+            return self.cuda_graph_runner.replay(batch)
+
+        input_metadata = InputMetadata.create(
+            self,
+            forward_mode=ForwardMode.DECODE,
+            req_pool_indices=batch.req_pool_indices,
+            seq_lens=batch.seq_lens,
+            prefix_lens=batch.prefix_lens,
+            position_ids_offsets=batch.position_ids_offsets,
+            out_cache_loc=batch.out_cache_loc,
+            top_logprobs_nums=batch.top_logprobs_nums,
+            return_logprob=batch.return_logprob,
+        )
+        return self.model.forward(
+            batch.input_ids, input_metadata.positions, input_metadata
+        )
+
+    @torch.inference_mode()
+    def forward_extend(self, batch: Batch):
+        input_metadata = InputMetadata.create(
+            self,
+            forward_mode=ForwardMode.EXTEND,
+            req_pool_indices=batch.req_pool_indices,
+            seq_lens=batch.seq_lens,
+            prefix_lens=batch.prefix_lens,
+            position_ids_offsets=batch.position_ids_offsets,
+            out_cache_loc=batch.out_cache_loc,
+            top_logprobs_nums=batch.top_logprobs_nums,
+            return_logprob=batch.return_logprob,
+        )
+        return self.model.forward(
+            batch.input_ids, input_metadata.positions, input_metadata
+        )
+
+    @torch.inference_mode()
+    def forward_extend_multi_modal(self, batch: Batch):
+        input_metadata = InputMetadata.create(
+            self,
+            forward_mode=ForwardMode.EXTEND,
+            req_pool_indices=batch.req_pool_indices,
+            seq_lens=batch.seq_lens,
+            prefix_lens=batch.prefix_lens,
+            position_ids_offsets=batch.position_ids_offsets,
+            out_cache_loc=batch.out_cache_loc,
+            return_logprob=batch.return_logprob,
+            top_logprobs_nums=batch.top_logprobs_nums,
+        )
+        return self.model.forward(
+            batch.input_ids,
+            input_metadata.positions,
+            input_metadata,
+            batch.pixel_values,
+            batch.image_sizes,
+            batch.image_offsets,
+        )
+
+    def forward(self, batch: Batch, forward_mode: ForwardMode):
+        if self.is_multimodal_model and forward_mode == ForwardMode.EXTEND:
+            return self.forward_extend_multi_modal(batch)
+        elif forward_mode == ForwardMode.DECODE:
+            return self.forward_decode(batch)
+        elif forward_mode == ForwardMode.EXTEND:
+            return self.forward_extend(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 tmp in entry:
+                        model_arch_name_to_cls[tmp.__name__] = tmp
+                else:
+                    model_arch_name_to_cls[entry.__name__] = entry
+
+            # compat: some models such as chatglm has incorrect class set in config.json
+            # usage: [ tuple("From_Entry_Class_Name": EntryClass), ]
+            if hasattr(module, "EntryClassRemapping") and isinstance(
+                module.EntryClassRemapping, list
+            ):
+                for remap in module.EntryClassRemapping:
+                    if isinstance(remap, tuple) and len(remap) == 2:
+                        model_arch_name_to_cls[remap[0]] = remap[1]
+
+    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 → controller}/radix_cache.py

@@ -1,8 +1,10 @@
+"""
+The radix tree data structure for managing the KV cache.
+"""
+
 import heapq
 import time
 from collections import defaultdict
-from dataclasses import dataclass
-from typing import Tuple
 
 import torch
 
@@ -11,34 +13,38 @@ class TreeNode:
     def __init__(self):
         self.children = defaultdict(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):
+    def __lt__(self, other: "TreeNode"):
         return self.last_access_time < other.last_access_time
 
 
-def match(key, seq):
+def _key_match(key0, key1):
     i = 0
-    for k, w in zip(key, seq):
-        if k != w:
+    for k0, k1 in zip(key0, key1):
+        if k0 != k1:
             break
         i += 1
     return i
 
 
 class RadixCache:
-    def __init__(self, disable=False):
-        self.reset()
+    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()
 
     ##### Public API #####
 
     def reset(self):
         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 +56,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.free(indices)
+            else:
+                return torch.tensor([], dtype=torch.int64), self.root_node
+
+        # Radix Cache takes one ref in memory pool
+        self.token_to_kv_pool.free(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()}")
@@ -69,7 +111,7 @@ class RadixCache:
 
     def evict(self, num_tokens, evict_callback):
         if self.disable:
-            raise RuntimeError()
+            return
 
         leaves = self._collect_leaves()
         heapq.heapify(leaves)
@@ -80,32 +122,33 @@ 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)
+            evict_callback(x.value)
+            num_evicted += len(x.value)
             self._delete_leaf(x)
 
             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
 
@@ -113,42 +156,48 @@ class RadixCache:
         return self.evictable_size_
 
     ##### Internal Helper Functions #####
+
     def _match_prefix_helper(self, node, key, value, last_node):
         node.last_access_time = time.time()
-
-        for c_key, child in node.children.items():
-            prefix_len = match(c_key, key)
-            if prefix_len != 0:
-                if prefix_len < len(c_key):
-                    new_node = self._split_node(c_key, child, prefix_len)
-                    value.append(new_node.value)
-                    last_node[0] = new_node
-                else:
-                    value.append(child.value)
-                    last_node[0] = child
-                    self._match_prefix_helper(child, key[prefix_len:], value, last_node)
-                break
-
-    def _split_node(self, key, child, split_len):
+        if len(key) == 0:
+            return
+
+        if key[0] in node.children.keys():
+            child = node.children[key[0]]
+            prefix_len = _key_match(child.key, key)
+            if prefix_len < len(child.key):
+                new_node = self._split_node(child.key, child, prefix_len)
+                value.append(new_node.value)
+                last_node[0] = new_node
+            else:
+                value.append(child.value)
+                last_node[0] = child
+                self._match_prefix_helper(child, key[prefix_len:], value, last_node)
+
+    def _split_node(self, key, child: TreeNode, split_len):
         # new_node -> child
         new_node = TreeNode()
-        new_node.children = {key[split_len:]: child}
+        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
+        child.key = child.key[split_len:]
         child.value = child.value[split_len:]
-        new_node.parent.children[key[:split_len]] = new_node
-        del new_node.parent.children[key]
+        new_node.parent.children[key[:split_len][0]] = new_node
         return new_node
 
     def _insert_helper(self, node, key, value):
         node.last_access_time = time.time()
+        if len(key) == 0:
+            return 0
 
-        for c_key, child in node.children.items():
-            prefix_len = match(c_key, key)
+        if key[0] in node.children.keys():
+            child = node.children[key[0]]
+            prefix_len = _key_match(child.key, key)
 
-            if prefix_len == len(c_key):
+            if prefix_len == len(child.key):
                 if prefix_len == len(key):
                     return prefix_len
                 else:
@@ -156,23 +205,23 @@ class RadixCache:
                     value = value[prefix_len:]
                     return prefix_len + self._insert_helper(child, key, value)
 
-            if prefix_len:
-                new_node = self._split_node(c_key, child, prefix_len)
-                return prefix_len + self._insert_helper(
-                    new_node, key[prefix_len:], value[prefix_len:]
-                )
+            new_node = self._split_node(child.key, child, prefix_len)
+            return prefix_len + self._insert_helper(
+                new_node, key[prefix_len:], value[prefix_len:]
+            )
 
         if len(key):
             new_node = TreeNode()
             new_node.parent = node
+            new_node.key = key
             new_node.value = value
-            node.children[key] = new_node
+            node.children[key[0]] = new_node
             self.evictable_size_ += len(value)
         return 0
 
-    def _print_helper(self, node, indent):
-        for key, child in node.children.items():
-            print(" " * indent, len(key), key[:10], f"r={child.ref_counter}")
+    def _print_helper(self, node: TreeNode, indent):
+        for _, child in node.children.items():
+            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):
@@ -180,7 +229,7 @@ class RadixCache:
             if v == node:
                 break
         del node.parent.children[k]
-        self.evictable_size_ -= len(k)
+        self.evictable_size_ -= len(node.key)
 
     def _total_size_helper(self, node):
         x = len(node.value)
@@ -203,7 +252,7 @@ class RadixCache:
 
 
 if __name__ == "__main__":
-    tree = RadixCache(disable=False)
+    tree = RadixCache(None, None, False)
 
     tree.insert("Hello")
     tree.insert("Hello")