sglang 0.5.2rc1__py3-none-any.whl → 0.5.3__py3-none-any.whl

sglang/srt/disaggregation/launch_lb.py

@@ -1,118 +0,0 @@
-import argparse
-import dataclasses
-
-from sglang.srt.disaggregation.mini_lb import PrefillConfig, run
-
-
-@dataclasses.dataclass
-class LBArgs:
-    host: str = "0.0.0.0"
-    port: int = 8000
-    policy: str = "random"
-    prefill_infos: list = dataclasses.field(default_factory=list)
-    decode_infos: list = dataclasses.field(default_factory=list)
-    log_interval: int = 5
-    timeout: int = 600
-
-    @staticmethod
-    def add_cli_args(parser: argparse.ArgumentParser):
-        parser.add_argument(
-            "--host",
-            type=str,
-            default=LBArgs.host,
-            help=f"Host to bind the server (default: {LBArgs.host})",
-        )
-        parser.add_argument(
-            "--port",
-            type=int,
-            default=LBArgs.port,
-            help=f"Port to bind the server (default: {LBArgs.port})",
-        )
-        parser.add_argument(
-            "--policy",
-            type=str,
-            default=LBArgs.policy,
-            choices=["random", "po2"],
-            help=f"Policy to use for load balancing (default: {LBArgs.policy})",
-        )
-        parser.add_argument(
-            "--prefill",
-            type=str,
-            default=[],
-            nargs="+",
-            help="URLs for prefill servers",
-        )
-        parser.add_argument(
-            "--decode",
-            type=str,
-            default=[],
-            nargs="+",
-            help="URLs for decode servers",
-        )
-        parser.add_argument(
-            "--prefill-bootstrap-ports",
-            type=int,
-            nargs="+",
-            help="Bootstrap ports for prefill servers",
-        )
-        parser.add_argument(
-            "--log-interval",
-            type=int,
-            default=LBArgs.log_interval,
-            help=f"Log interval in seconds (default: {LBArgs.log_interval})",
-        )
-        parser.add_argument(
-            "--timeout",
-            type=int,
-            default=LBArgs.timeout,
-            help=f"Timeout in seconds (default: {LBArgs.timeout})",
-        )
-
-    @classmethod
-    def from_cli_args(cls, args: argparse.Namespace) -> "LBArgs":
-        bootstrap_ports = args.prefill_bootstrap_ports
-        if bootstrap_ports is None:
-            bootstrap_ports = [None] * len(args.prefill)
-        elif len(bootstrap_ports) == 1:
-            bootstrap_ports = bootstrap_ports * len(args.prefill)
-        else:
-            if len(bootstrap_ports) != len(args.prefill):
-                raise ValueError(
-                    "Number of prefill URLs must match number of bootstrap ports"
-                )
-
-        prefill_infos = [
-            (url, port) for url, port in zip(args.prefill, bootstrap_ports)
-        ]
-
-        return cls(
-            host=args.host,
-            port=args.port,
-            policy=args.policy,
-            prefill_infos=prefill_infos,
-            decode_infos=args.decode,
-            log_interval=args.log_interval,
-            timeout=args.timeout,
-        )
-
-
-def main():
-    parser = argparse.ArgumentParser(
-        description="PD Disaggregation Load Balancer Server"
-    )
-    LBArgs.add_cli_args(parser)
-    args = parser.parse_args()
-    lb_args = LBArgs.from_cli_args(args)
-
-    prefill_configs = [PrefillConfig(url, port) for url, port in lb_args.prefill_infos]
-    run(
-        prefill_configs,
-        lb_args.decode_infos,
-        lb_args.host,
-        lb_args.port,
-        lb_args.timeout,
-    )
-
-
-if __name__ == "__main__":
-    main()

sglang/srt/mem_cache/lora_radix_cache.py

@@ -1,421 +0,0 @@
-"""Radix cache for LoRA. It's modified based on RadixCache with lora_id added to the key of nodes."""
-
-import heapq
-import time
-from collections import defaultdict
-from typing import TYPE_CHECKING, Any, List, Optional
-
-import torch
-
-from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
-from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
-from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
-
-if TYPE_CHECKING:
-    from sglang.srt.managers.schedule_batch import Req
-else:
-    Req = Any  # Placeholder for Req type when not type checking
-
-
-class LoRAKey:
-
-    def __init__(self, lora_id: str, token_ids: List[int]):
-        self.lora_id = (
-            lora_id  # lora_id of adaptor, should be hash value of adaptor path
-        )
-        self.token_ids = token_ids  # token_ids of the key
-
-    def __len__(self):
-        return len(self.token_ids)
-
-
-def get_child_key(key: LoRAKey):
-    # Here the key of children dict is the hash of lora_id + str(token_ids[0])
-    # So the child key can be matched only when lora_id and token_ids[0] are the same
-    if key.lora_id is None:
-        return hash(str(key.token_ids[0]))
-    else:
-        return hash(key.lora_id + str(key.token_ids[0]))
-
-
-class LoRATreeNode:
-
-    counter = 0
-
-    def __init__(self, id: Optional[int] = None):
-        self.children = defaultdict(LoRATreeNode)
-        self.parent: LoRATreeNode = None
-        self.key: LoRAKey = None
-        self.value: Optional[torch.Tensor] = None
-        self.lock_ref = 0
-        self.last_access_time = time.monotonic()
-
-        self.id = LoRATreeNode.counter if id is None else id
-        LoRATreeNode.counter += 1
-
-    @property
-    def evicted(self):
-        return self.value is None
-
-    def __lt__(self, other: "LoRATreeNode"):
-        return self.last_access_time < other.last_access_time
-
-
-def _key_match(key0: LoRAKey, key1: LoRAKey):
-    if key0.lora_id != key1.lora_id:
-        raise ValueError(
-            f"_key_match should be run on the same lora_id, but got key0.lora_id={key0.lora_id} != key1.lora_id={key1.lora_id}"
-        )
-    i = 0
-    for k0, k1 in zip(key0.token_ids, key1.token_ids):
-        if k0 != k1:
-            break
-        i += 1
-    return i
-
-
-class LoRARadixCache(BasePrefixCache):
-
-    def __init__(
-        self,
-        req_to_token_pool: ReqToTokenPool,
-        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
-        page_size: int,
-        disable: bool = False,
-    ):
-        if page_size > 1:
-            raise ValueError("LoRARadixCache currently only supports page_size = 1")
-
-        if token_to_kv_pool_allocator is None:
-            raise ValueError(
-                "token_to_kv_pool_allocator is required to run LoraRadixCache"
-            )
-
-        self.req_to_token_pool = req_to_token_pool
-        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
-        self.page_size = page_size
-        self.disable = disable
-        self.device = self.token_to_kv_pool_allocator.device
-
-        self.key_match_fn = _key_match
-        self.get_child_key_fn = get_child_key
-        self.reset()
-
-    def reset(self):
-        self.root_node = LoRATreeNode()
-        self.root_node.key = LoRAKey(lora_id="", token_ids=[])
-        self.root_node.value = None
-        self.evictable_size_ = 0
-        self.protected_size_ = 0
-
-    def match_prefix(self, key: List[int], **kwargs) -> MatchResult:
-        raise ValueError(
-            "LoRARadixCache needs both token ids and lora id as inputs for matching. Please use match_prefix_with_lora_id instead."
-        )
-
-    def match_prefix_with_lora_id(self, key: LoRAKey, **kwargs) -> MatchResult:
-        """Find the matching prefix from the lora radix tree.
-        Args:
-            key: A LoRAKey to find a matching prefix.
-        Returns:
-            A tuple of a tensor of matching prefix token IDs and
-            the last node that contains the prefix values. Note that
-            this API can modify the internal state of the Radix tree.
-            The last node create a new child if the prefix is shorter
-            than the last node's value.
-        """
-        if self.disable or len(key) == 0:
-            return MatchResult(
-                device_indices=torch.empty(
-                    (0,),
-                    dtype=torch.int64,
-                    device=self.device,
-                ),
-                last_device_node=self.root_node,
-                last_host_node=self.root_node,
-            )
-
-        value, last_node = self._match_prefix_helper(self.root_node, key)
-        if value:
-            value = torch.cat(value)
-        else:
-            value = torch.empty((0,), dtype=torch.int64, device=self.device)
-        return MatchResult(
-            device_indices=value,
-            last_device_node=last_node,
-            last_host_node=last_node,
-        )
-
-    def insert(self, key: LoRAKey, value=None):
-        if self.disable:
-            return 0
-
-        if value is None:
-            value = [x for x in key.token_ids]
-        return self._insert_helper(self.root_node, key, value)
-
-    def cache_finished_req(self, req: Req):
-        """Cache request when it finishes."""
-        if self.disable:
-            kv_indices = self.req_to_token_pool.req_to_token[
-                req.req_pool_idx, : len(req.origin_input_ids) + len(req.output_ids) - 1
-            ]
-            self.token_to_kv_pool_allocator.free(kv_indices)
-            self.req_to_token_pool.free(req.req_pool_idx)
-            return
-
-        token_ids = (req.origin_input_ids + req.output_ids)[:-1]
-        kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, : len(token_ids)
-        ]
-
-        page_aligned_len = len(kv_indices)
-        page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True)
-
-        # Radix Cache takes one ref in memory pool
-        lora_key = LoRAKey(lora_id=req.lora_id, token_ids=token_ids[:page_aligned_len])
-        new_prefix_len = self.insert(lora_key, page_aligned_kv_indices)
-        self.token_to_kv_pool_allocator.free(
-            kv_indices[len(req.prefix_indices) : new_prefix_len]
-        )
-
-        # Remove req slot release the cache lock
-        self.req_to_token_pool.free(req.req_pool_idx)
-        self.dec_lock_ref(req.last_node)
-
-    def cache_unfinished_req(self, req: Req, chunked=False):
-        """Cache request when it is unfinished."""
-        if self.disable:
-            return
-
-        token_ids = req.fill_ids
-        kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, : len(token_ids)
-        ]
-
-        page_aligned_len = len(kv_indices)
-        page_aligned_kv_indices = kv_indices.to(dtype=torch.int64, copy=True)
-        page_aligned_token_ids = token_ids[:page_aligned_len]
-
-        # Radix Cache takes one ref in memory pool
-        inserted_key = LoRAKey(lora_id=req.lora_id, token_ids=page_aligned_token_ids)
-        new_prefix_len = self.insert(inserted_key, page_aligned_kv_indices)
-        self.token_to_kv_pool_allocator.free(
-            kv_indices[len(req.prefix_indices) : new_prefix_len]
-        )
-
-        # The prefix indices could be updated, reuse it
-        new_indices, new_last_node, _, _ = self.match_prefix_with_lora_id(inserted_key)
-        self.req_to_token_pool.write(
-            (req.req_pool_idx, slice(len(req.prefix_indices), len(new_indices))),
-            new_indices[len(req.prefix_indices) :],
-        )
-
-        self.dec_lock_ref(req.last_node)
-        self.inc_lock_ref(new_last_node)
-
-        # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
-        req.prefix_indices = new_indices
-        req.last_node = new_last_node
-
-    def pretty_print(self):
-        self._print_helper(self.root_node, 0)
-        print(f"#tokens: {self.total_size()}")
-
-    def total_size(self):
-        return self._total_size_helper()
-
-    def evict(self, num_tokens: int):
-        if self.disable:
-            return
-
-        leaves = self._collect_leaves()
-        heapq.heapify(leaves)
-
-        num_evicted = 0
-        while num_evicted < num_tokens and len(leaves):
-            x = heapq.heappop(leaves)
-
-            if x == self.root_node:
-                break
-            if x.lock_ref > 0:
-                continue
-
-            self.token_to_kv_pool_allocator.free(x.value)
-            num_evicted += len(x.value)
-            self._delete_leaf(x)
-
-            if len(x.parent.children) == 0:
-                heapq.heappush(leaves, x.parent)
-
-    def inc_lock_ref(self, node: LoRATreeNode):
-        if self.disable:
-            return 0
-
-        delta = 0
-        while node != self.root_node:
-            if node.lock_ref == 0:
-                self.evictable_size_ -= len(node.value)
-                self.protected_size_ += len(node.value)
-                delta -= len(node.value)
-            node.lock_ref += 1
-            node = node.parent
-        return delta
-
-    def dec_lock_ref(self, node: LoRATreeNode):
-        if self.disable:
-            return 0
-
-        delta = 0
-        while node != self.root_node:
-            if node.lock_ref == 1:
-                self.evictable_size_ += len(node.value)
-                self.protected_size_ -= len(node.value)
-                delta += len(node.value)
-            node.lock_ref -= 1
-            node = node.parent
-        return delta
-
-    def evictable_size(self):
-        return self.evictable_size_
-
-    def protected_size(self):
-        # protected size refers to the size of the cache that is locked
-        return self.protected_size_
-
-    def all_values_flatten(self):
-        values = []
-
-        def _dfs_helper(node: LoRATreeNode):
-            for _, child in node.children.items():
-                values.append(child.value)
-                _dfs_helper(child)
-
-        _dfs_helper(self.root_node)
-        return torch.cat(values)
-
-    ##### Internal Helper Functions #####
-
-    def _match_prefix_helper(self, node: LoRATreeNode, key: LoRAKey):
-        node.last_access_time = time.monotonic()
-
-        child_key = self.get_child_key_fn(key)
-
-        value = []
-        while len(key) > 0 and child_key in node.children.keys():
-            child = node.children[child_key]
-            child.last_access_time = time.monotonic()
-            prefix_len = self.key_match_fn(child.key, key)
-            if prefix_len < len(child.key):
-                new_node = self._split_node(child.key, child, prefix_len)
-                value.append(new_node.value)
-                node = new_node
-                break
-            else:
-                value.append(child.value)
-                node = child
-                key = LoRAKey(lora_id=key.lora_id, token_ids=key.token_ids[prefix_len:])
-
-                if len(key):
-                    child_key = self.get_child_key_fn(key)
-
-        return value, node
-
-    def _split_node(self, key: LoRAKey, child: LoRATreeNode, split_len: int):
-        # new_node -> child
-        new_node = LoRATreeNode()
-        key_split_1 = LoRAKey(lora_id=key.lora_id, token_ids=key.token_ids[:split_len])
-        key_split_2 = LoRAKey(lora_id=key.lora_id, token_ids=key.token_ids[split_len:])
-        new_node.children = {self.get_child_key_fn(key_split_2): child}
-        new_node.parent = child.parent
-        new_node.lock_ref = child.lock_ref
-        new_node.key = key_split_1
-        new_node.value = child.value[:split_len]
-        child.parent = new_node
-        child.key = key_split_2
-        child.value = child.value[split_len:]
-        new_node.parent.children[self.get_child_key_fn(key)] = new_node
-
-        return new_node
-
-    def _insert_helper(self, node: LoRATreeNode, key: LoRAKey, value):
-        node.last_access_time = time.monotonic()
-        if len(key) == 0:
-            return 0
-
-        child_key = self.get_child_key_fn(key)
-
-        total_prefix_length = 0
-        while len(key) > 0 and child_key in node.children.keys():
-            node = node.children[child_key]
-            node.last_access_time = time.monotonic()
-            prefix_len = self.key_match_fn(node.key, key)
-            total_prefix_length += prefix_len
-            key = LoRAKey(lora_id=key.lora_id, token_ids=key.token_ids[prefix_len:])
-            value = value[prefix_len:]
-
-            if prefix_len < len(node.key):
-                new_node = self._split_node(node.key, node, prefix_len)
-                node = new_node
-
-            if len(key):
-                child_key = self.get_child_key_fn(key)
-
-        if len(key):
-            new_node = LoRATreeNode()
-            new_node.parent = node
-            new_node.key = key
-            new_node.value = value
-            node.children[child_key] = new_node
-            self.evictable_size_ += len(value)
-        return total_prefix_length
-
-    def _print_helper(self, node: LoRATreeNode, indent: int):
-        """Prints the radix tree in a human-readable format."""
-        stack = [(node, indent)]
-        while stack:
-            current_node, current_indent = stack.pop()
-            print(
-                " " * current_indent,
-                len(current_node.key),
-                current_node.key.token_ids[:10],
-                f"r={current_node.lock_ref}",
-            )
-            for key, child in current_node.children.items():
-                stack.append((child, current_indent + 2))
-
-                assert key == self.get_child_key_fn(
-                    child.key
-                ), f"{key=}, {self.get_child_key_fn(child.key)=}"
-
-    def _delete_leaf(self, node):
-        for k, v in node.parent.children.items():
-            if v == node:
-                break
-        del node.parent.children[k]
-        self.evictable_size_ -= len(node.key)
-
-    def _total_size_helper(self):
-        total_size = 0
-        stack = [self.root_node]
-        while stack:
-            current_node = stack.pop()
-            total_size += len(current_node.value)
-            for child in current_node.children.values():
-                if child.evicted:
-                    continue
-                stack.append(child)
-        return total_size
-
-    def _collect_leaves(self):
-        ret_list = []
-        stack = [self.root_node]
-
-        while stack:
-            cur_node = stack.pop()
-            if len(cur_node.children) == 0:
-                ret_list.append(cur_node)
-            else:
-                stack.extend(cur_node.children.values())
-
-        return ret_list

sglang/srt/mem_cache/storage/mooncake_store/unit_test.py

@@ -1,40 +0,0 @@
-import torch
-from mooncake_store import MooncakeStore
-
-
-def test_init_and_warmup():
-    store = MooncakeStore()
-    assert store.store is not None
-
-
-def test_register_buffer():
-    store = MooncakeStore()
-    tensor = torch.zeros(1024, dtype=torch.float32)
-    store.register_buffer(tensor)
-
-
-def test_set_and_get():
-    store = MooncakeStore()
-
-    key = ["test_key_" + str(i) for i in range(2)]
-    tensor = torch.arange(256, dtype=torch.float32).cuda()
-    ptrs = [tensor.data_ptr(), tensor.data_ptr()]
-    sizes = [tensor.numel() * tensor.element_size()] * 2
-
-    store.set(key, target_location=ptrs, target_sizes=sizes)
-    store.get(key, target_location=ptrs, target_sizes=sizes)
-
-
-def test_exists():
-    store = MooncakeStore()
-    keys = ["test_key_0", "non_existent_key"]
-    result = store.exists(keys)
-    assert isinstance(result, dict)
-    assert "test_key_0" in result
-
-
-if __name__ == "__main__":
-    test_init_and_warmup()
-    test_register_buffer()
-    test_set_and_get()
-    test_exists()