sglang 0.4.3.post2__py3-none-any.whl → 0.4.3.post4__py3-none-any.whl

sglang/srt/managers/tp_worker_overlap_thread.py

@@ -100,7 +100,7 @@ class TpModelWorkerClient:
     def get_memory_pool(self):
         return (
             self.worker.model_runner.req_to_token_pool,
-            self.worker.model_runner.token_to_kv_pool,
+            self.worker.model_runner.token_to_kv_pool_allocator,
         )
 
     def forward_thread_func(self):
@@ -175,7 +175,7 @@ class TpModelWorkerClient:
                 logits_output.next_token_logprobs.tolist()
             )
             if logits_output.input_token_logprobs is not None:
-                logits_output.input_token_logprobs = (
+                logits_output.input_token_logprobs = tuple(
                     logits_output.input_token_logprobs.tolist()
                 )
         next_token_ids = next_token_ids.tolist()
@@ -188,8 +188,7 @@ class TpModelWorkerClient:
         model_worker_batch.sampling_info = self.cur_sampling_info = dataclasses.replace(
             sampling_info,
             sampling_info_done=threading.Event(),
-            scaling_penalties=sampling_info.scaling_penalties,
-            linear_penalties=sampling_info.linear_penalties,
+            penalizer_orchestrator=None,
         )
 
         # A cuda stream sync here to avoid the cuda illegal memory access error.

sglang/srt/mem_cache/chunk_cache.py

@@ -1,29 +1,33 @@
 from __future__ import annotations
 
 """Cache for chunked prefill, used when RadixCache is disabled."""
+from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple
 
-from typing import TYPE_CHECKING, Callable, List, Optional, Tuple
+import torch
 
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
-from sglang.srt.mem_cache.memory_pool import BaseTokenToKVPool, ReqToTokenPool
+from sglang.srt.mem_cache.memory_pool import ReqToTokenPool, TokenToKVPoolAllocator
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
 
 
 class ChunkCacheEntry:
-    def __init__(self, rid, value):
+    def __init__(self, rid: str, value: torch.Tensor):
         self.rid = rid
         self.value = value
 
 
 class ChunkCache(BasePrefixCache):
     def __init__(
-        self, req_to_token_pool: ReqToTokenPool, token_to_kv_pool: BaseTokenToKVPool
+        self,
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool_allocator: TokenToKVPoolAllocator,
     ):
         self.disable = True
         self.req_to_token_pool = req_to_token_pool
-        self.token_to_kv_pool = token_to_kv_pool
+        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
+        self.entries: Dict[str, ChunkCacheEntry] = {}
 
         self.reset()
 
@@ -48,16 +52,13 @@ class ChunkCache(BasePrefixCache):
             req.req_pool_idx, :token_id_len
         ]
         self.req_to_token_pool.free(req.req_pool_idx)
-        self.token_to_kv_pool.free(kv_indices)
+        self.token_to_kv_pool_allocator.free(kv_indices)
 
         if req.rid in self.entries:
             del self.entries[req.rid]
 
-    def cache_unfinished_req(self, req: Req, token_ids: Optional[List[int]] = None):
-        if token_ids is None:
-            token_id_len = len(req.fill_ids)
-        else:
-            token_id_len = len(token_ids)
+    def cache_unfinished_req(self, req: Req):
+        token_id_len = len(req.fill_ids)
 
         kv_indices = self.req_to_token_pool.req_to_token[
             req.req_pool_idx, :token_id_len
@@ -86,5 +87,11 @@ class ChunkCache(BasePrefixCache):
     def evictable_size(self):
         return 0
 
+    def pretty_print(self):
+        return ""
+
     def protected_size(self):
         return 0
+
+    def pretty_print(self):
+        return ""

sglang/srt/mem_cache/hiradix_cache.py

@@ -0,0 +1,394 @@
+import heapq
+import logging
+import time
+from typing import List, Optional
+
+import torch
+
+from sglang.srt.managers.cache_controller import HiCacheController
+from sglang.srt.mem_cache.memory_pool import (
+    MHATokenToKVPool,
+    MHATokenToKVPoolHost,
+    ReqToTokenPool,
+)
+from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode, _key_match
+
+logger = logging.getLogger(__name__)
+
+
+class HiRadixCache(RadixCache):
+
+    def __init__(
+        self,
+        req_to_token_pool: ReqToTokenPool,
+        token_to_kv_pool: MHATokenToKVPool,
+    ):
+        self.token_to_kv_pool_host = MHATokenToKVPoolHost(token_to_kv_pool)
+        self.cache_controller = HiCacheController(
+            token_to_kv_pool, self.token_to_kv_pool_host
+        )
+
+        # record the nodes with ongoing write through
+        self.ongoing_write_through = {}
+        # record the node segments with ongoing load back
+        self.ongoing_load_back = {}
+        # todo: dynamically adjust the threshold
+        self.write_through_threshold = 1
+        self.load_back_threshold = 10
+        super().__init__(req_to_token_pool, token_to_kv_pool, disable=False)
+
+    def reset(self):
+        TreeNode.counter = 0
+        self.cache_controller.reset()
+        self.token_to_kv_pool_host.clear()
+        super().reset()
+
+    def get_height(self, node: TreeNode):
+        height = 0
+        while node != self.root_node:
+            node = node.parent
+            height += 1
+        return height
+
+    def write_backup(self, node: TreeNode):
+        host_indices = self.cache_controller.write(
+            device_indices=node.value,
+            priority=-self.get_height(node),
+            node_id=node.id,
+        )
+        if host_indices is None:
+            self.evict_host(len(node.value))
+            host_indices = self.cache_controller.write(
+                device_indices=node.value,
+                priority=-self.get_height(node),
+                node_id=node.id,
+            )
+        if host_indices is not None:
+            node.host_value = host_indices
+            self.ongoing_write_through[node.id] = node
+            self.inc_lock_ref(node)
+        else:
+            return None
+
+        return len(host_indices)
+
+    def inc_hit_count(self, node: TreeNode):
+        if self.cache_controller.write_policy != "write_through_selective":
+            return
+        node.hit_count += 1
+        if node.host_value is None and node.hit_count > self.write_through_threshold:
+            self.write_backup(node)
+            node.hit_count = 0
+
+    def writing_check(self):
+        while not self.cache_controller.ack_write_queue.empty():
+            try:
+                ack_id = self.cache_controller.ack_write_queue.get_nowait()
+                self.dec_lock_ref(self.ongoing_write_through[ack_id])
+                # clear the reference
+                del self.ongoing_write_through[ack_id]
+            except Exception:
+                break
+
+    def loading_check(self):
+        while not self.cache_controller.ack_load_queue.empty():
+            try:
+                ack_id = self.cache_controller.ack_load_queue.get_nowait()
+                start_node, end_node = self.ongoing_load_back[ack_id]
+                self.dec_lock_ref(end_node)
+                while end_node != start_node:
+                    assert end_node.loading
+                    end_node.loading = False
+                    end_node = end_node.parent
+                # clear the reference
+                del self.ongoing_load_back[ack_id]
+            except Exception:
+                break
+
+    def evictable_size(self):
+        self.writing_check()
+        self.loading_check()
+        return self.evictable_size_
+
+    def evict(self, num_tokens: int, evict_callback=None):
+        leaves = self._collect_leaves_device()
+        heapq.heapify(leaves)
+
+        num_evicted = 0
+        pending_nodes = []
+        while num_evicted < num_tokens and len(leaves):
+            x = heapq.heappop(leaves)
+
+            if x.lock_ref > 0:
+                continue
+
+            if x.host_value is None:
+                if self.cache_controller.write_policy == "write_back":
+                    num_evicted += self.write_backup(x)
+                elif self.cache_controller.write_policy == "write_through_selective":
+                    num_evicted += self._evict_write_through_selective(x)
+                else:
+                    assert (
+                        self.cache_controller.write_policy != "write_through"
+                    ), "write_through should be inclusive"
+                    raise NotImplementedError
+            else:
+                num_evicted += self._evict_write_through(x)
+
+            for child in x.parent.children.values():
+                if child in pending_nodes:
+                    continue
+                if not child.evicted:
+                    break
+            else:
+                # all children are evicted or no children
+                heapq.heappush(leaves, x.parent)
+
+        if self.cache_controller.write_policy == "write_back":
+            # blocking till all write back complete
+            while len(self.ongoing_write_through) > 0:
+                self.writing_check()
+                time.sleep(0.1)
+
+    def _evict_write_through(self, node: TreeNode):
+        # evict a node already written to host
+        num_evicted = self.cache_controller.evict_device(node.value, node.host_value)
+        assert num_evicted > 0
+        self.evictable_size_ -= num_evicted
+        node.value = None
+        return num_evicted
+
+    def _evict_write_through_selective(self, node: TreeNode):
+        # evict a node not initiated write to host
+        self.cache_controller.mem_pool_device.free(node.value)
+        num_evicted = len(node.value)
+        self._delete_leaf(node)
+        return num_evicted
+
+    def evict_host(self, num_tokens: int):
+        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
+            # only evict the host value of evicted nodes
+            if not x.evicted:
+                continue
+            assert x.lock_ref == 0 and x.host_value is not None
+
+            assert self.cache_controller.evict_host(x.host_value) > 0
+            for k, v in x.parent.children.items():
+                if v == x:
+                    break
+            del x.parent.children[k]
+
+            if len(x.parent.children) == 0 and x.parent.evicted:
+                heapq.heappush(leaves, x.parent)
+
+    def load_back(
+        self, node: TreeNode, mem_quota: Optional[int] = None
+    ) -> Optional[torch.Tensor]:
+        # todo: more loading policies
+
+        last_hit_node = node
+        nodes_to_load = []
+        while node.evicted:
+            assert (
+                node.backuped
+            ), "No backup available on evicted nodes, should not happen"
+            nodes_to_load.insert(0, node)
+            node = node.parent
+        else:
+            ancester_node = node
+
+        # protect the ancestor nodes from eviction
+        delta = self.inc_lock_ref(ancester_node)
+
+        # load it all or not at all
+        host_indices = torch.cat([n.host_value for n in nodes_to_load])
+        if len(host_indices) < self.load_back_threshold or (
+            len(host_indices) > mem_quota + delta if mem_quota is not None else False
+        ):
+            # skip loading back if the total size is too small or exceeding the memory quota
+            self.dec_lock_ref(ancester_node)
+            return None
+
+        device_indices = self.cache_controller.load(
+            host_indices=host_indices, node_id=last_hit_node.id
+        )
+        if device_indices is None:
+            self.evict(len(host_indices))
+            device_indices = self.cache_controller.load(
+                host_indices=host_indices, node_id=last_hit_node.id
+            )
+        self.dec_lock_ref(ancester_node)
+        if device_indices is None:
+            # no sufficient GPU memory to load back KV caches
+            return None
+
+        self.ongoing_load_back[last_hit_node.id] = (ancester_node, last_hit_node)
+        offset = 0
+        for node in nodes_to_load:
+            node.value = device_indices[offset : offset + len(node.host_value)]
+            offset += len(node.host_value)
+            node.loading = True
+        self.evictable_size_ += len(device_indices)
+        self.inc_lock_ref(last_hit_node)
+
+        return device_indices
+
+    def loading_complete(self, node: TreeNode):
+        self.loading_check()
+        return node.loading == False
+
+    def init_load_back(
+        self,
+        last_node: TreeNode,
+        prefix_indices: torch.Tensor,
+        mem_quota: Optional[int] = None,
+    ):
+        assert (
+            len(prefix_indices) == 0 or prefix_indices.is_cuda
+        ), "indices of device kV caches should be on GPU"
+        if last_node.evicted:
+            loading_values = self.load_back(last_node, mem_quota)
+            if loading_values is not None:
+                prefix_indices = (
+                    loading_values
+                    if len(prefix_indices) == 0
+                    else torch.cat([prefix_indices, loading_values])
+                )
+                logger.debug(
+                    f"loading back {len(loading_values)} tokens for node {last_node.id}"
+                )
+
+            while last_node.evicted:
+                last_node = last_node.parent
+
+        return last_node, prefix_indices
+
+    def _match_prefix_helper(
+        self, node: TreeNode, key: List, value, last_node: TreeNode
+    ):
+        node.last_access_time = time.time()
+        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)
+                self.inc_hit_count(new_node)
+                if not new_node.evicted:
+                    value.append(new_node.value)
+                last_node[0] = new_node
+            else:
+                self.inc_hit_count(child)
+                if not child.evicted:
+                    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: int):
+        # child node split into new_node -> child
+        new_node = TreeNode()
+        new_node.children = {key[split_len]: child}
+        new_node.parent = child.parent
+        new_node.lock_ref = child.lock_ref
+        new_node.key = child.key[:split_len]
+        new_node.loading = child.loading
+
+        # split value and host value if exists
+        if child.evicted:
+            new_node.value = None
+        else:
+            new_node.value = child.value[:split_len]
+            child.value = child.value[split_len:]
+        if child.host_value is not None:
+            new_node.host_value = child.host_value[:split_len]
+            child.host_value = child.host_value[split_len:]
+        child.parent = new_node
+        child.key = child.key[split_len:]
+        new_node.parent.children[key[0]] = new_node
+        return new_node
+
+    def _insert_helper(self, node: TreeNode, key: List, value):
+        node.last_access_time = time.time()
+        if len(key) == 0:
+            return 0
+
+        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):
+                if child.evicted:
+                    # change the reference if the node is evicted
+                    # this often happens in the case of KV cache recomputation
+                    child.value = value[:prefix_len]
+                    self.token_to_kv_pool_host.update_synced(child.host_value)
+                    self.evictable_size_ += len(value[:prefix_len])
+                    return self._insert_helper(
+                        child, key[prefix_len:], value[prefix_len:]
+                    )
+                else:
+                    self.inc_hit_count(child)
+                    return prefix_len + self._insert_helper(
+                        child, key[prefix_len:], value[prefix_len:]
+                    )
+
+            # partial match, split the node
+            new_node = self._split_node(child.key, child, prefix_len)
+            if new_node.evicted:
+                new_node.value = value[:prefix_len]
+                self.token_to_kv_pool_host.update_synced(new_node.host_value)
+                self.evictable_size_ += len(new_node.value)
+                return self._insert_helper(
+                    new_node, key[prefix_len:], value[prefix_len:]
+                )
+            else:
+                self.inc_hit_count(new_node)
+                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[0]] = new_node
+            self.evictable_size_ += len(value)
+
+            if self.cache_controller.write_policy == "write_through":
+                self.write_backup(new_node)
+        return 0
+
+    def _collect_leaves_device(self):
+        def is_leaf(node):
+            if node.evicted:
+                return False
+            if node == self.root_node:
+                return False
+            if len(node.children) == 0:
+                return True
+            for child in node.children.values():
+                if not child.evicted:
+                    return False
+            return True
+
+        ret_list = []
+        stack = [self.root_node]
+        while stack:
+            cur_node = stack.pop()
+            if is_leaf(cur_node):
+                ret_list.append(cur_node)
+            else:
+                for cur_child in cur_node.children.values():
+                    if not cur_child.evicted:
+                        stack.append(cur_child)
+        return ret_list

sglang/srt/mem_cache/memory_pool.py

@@ -20,9 +20,11 @@ Memory pool.
 
 SGLang has two levels of memory pool.
 ReqToTokenPool maps a a request to its token locations.
-BaseTokenToKVPool maps a token location to its KV cache data.
+TokenToKVPoolAllocator manages the indices to kv cache data.
+KVCache actually holds the physical kv cache.
 """
 
+import abc
 import logging
 import threading
 from enum import IntEnum
@@ -89,22 +91,43 @@ class ReqToTokenPool:
         self.free_slots = list(range(self.size))
 
 
-class BaseTokenToKVPool:
-    """A memory pool that maps a token location to its kv cache data."""
+class KVCache(abc.ABC):
+
+    @abc.abstractmethod
+    def get_key_buffer(self, layer_id: int) -> torch.Tensor:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_value_buffer(self, layer_id: int) -> torch.Tensor:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def set_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+    ) -> None:
+        raise NotImplementedError()
+
+
+class TokenToKVPoolAllocator:
+    """An allocator managing the indices to kv cache data."""
 
     def __init__(
         self,
         size: int,
         dtype: torch.dtype,
         device: str,
+        kvcache: KVCache,
     ):
         self.size = size
         self.dtype = dtype
-        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
-            # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2
-            self.store_dtype = torch.uint8
-        else:
-            self.store_dtype = dtype
         self.device = device
 
         self.free_slots = None
@@ -112,9 +135,14 @@ class BaseTokenToKVPool:
         self.free_group = []
         self.clear()
 
+        self._kvcache = kvcache
+
     def available_size(self):
         return len(self.free_slots)
 
+    def get_kvcache(self):
+        return self._kvcache
+
     def alloc(self, need_size: int):
         if need_size > len(self.free_slots):
             return None
@@ -148,26 +176,8 @@ class BaseTokenToKVPool:
         self.is_in_free_group = False
         self.free_group = []
 
-    def get_key_buffer(self, layer_id: int) -> torch.Tensor:
-        raise NotImplementedError()
-
-    def get_value_buffer(self, layer_id: int) -> torch.Tensor:
-        raise NotImplementedError()
 
-    def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
-        raise NotImplementedError()
-
-    def set_kv_buffer(
-        self,
-        layer: RadixAttention,
-        loc: torch.Tensor,
-        cache_k: torch.Tensor,
-        cache_v: torch.Tensor,
-    ) -> None:
-        raise NotImplementedError()
-
-
-class MHATokenToKVPool(BaseTokenToKVPool):
+class MHATokenToKVPool(KVCache):
 
     def __init__(
         self,
@@ -179,8 +189,14 @@ class MHATokenToKVPool(BaseTokenToKVPool):
         device: str,
         enable_memory_saver: bool,
     ):
-        super().__init__(size, dtype, device)
-
+        self.size = size
+        self.dtype = dtype
+        self.device = device
+        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
+            # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2
+            self.store_dtype = torch.uint8
+        else:
+            self.store_dtype = dtype
         self.memory_saver_adapter = TorchMemorySaverAdapter.create(
             enable=enable_memory_saver
         )
@@ -192,7 +208,7 @@ class MHATokenToKVPool(BaseTokenToKVPool):
 
         k_size, v_size = self.get_kv_size_bytes()
         logger.info(
-            f"KV Cache is allocated. K size: {k_size / GB:.2f} GB, V size: {v_size / GB:.2f} GB."
+            f"KV Cache is allocated. #tokens: {size}, K size: {k_size / GB:.2f} GB, V size: {v_size / GB:.2f} GB"
         )
 
     def _create_buffers(self):
@@ -297,7 +313,7 @@ def copy_two_array(loc, dst_1, src_1, dst_2, src_2, dtype, store_dtype):
     dst_2[loc] = src_2.to(dtype).view(store_dtype)
 
 
-class MLATokenToKVPool(BaseTokenToKVPool):
+class MLATokenToKVPool(KVCache):
     def __init__(
         self,
         size: int,
@@ -308,8 +324,14 @@ class MLATokenToKVPool(BaseTokenToKVPool):
         device: str,
         enable_memory_saver: bool,
     ):
-        super().__init__(size, dtype, device)
-
+        self.size = size
+        self.dtype = dtype
+        self.device = device
+        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
+            # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2
+            self.store_dtype = torch.uint8
+        else:
+            self.store_dtype = dtype
         self.kv_lora_rank = kv_lora_rank
 
         memory_saver_adapter = TorchMemorySaverAdapter.create(
@@ -356,7 +378,7 @@ class MLATokenToKVPool(BaseTokenToKVPool):
             self.kv_buffer[layer_id][loc] = cache_k
 
 
-class DoubleSparseTokenToKVPool(BaseTokenToKVPool):
+class DoubleSparseTokenToKVPool(KVCache):
     def __init__(
         self,
         size: int,
@@ -368,8 +390,14 @@ class DoubleSparseTokenToKVPool(BaseTokenToKVPool):
         heavy_channel_num: int,
         enable_memory_saver: bool,
     ):
-        super().__init__(size, dtype, device)
-
+        self.size = size
+        self.dtype = dtype
+        self.device = device
+        if dtype in (torch.float8_e5m2, torch.float8_e4m3fn):
+            # NOTE: Store as torch.uint8 because Tensor.index_put is not implemented for torch.float8_e5m2
+            self.store_dtype = torch.uint8
+        else:
+            self.store_dtype = dtype
         memory_saver_adapter = TorchMemorySaverAdapter.create(
             enable=enable_memory_saver
         )
@@ -437,7 +465,7 @@ def synchronized(func):
     return wrapper
 
 
-class MLATokenToKVPoolHost:
+class MHATokenToKVPoolHost:
 
     def __init__(
         self,
@@ -502,6 +530,9 @@ class MLATokenToKVPoolHost:
     def get_flat_data(self, indices):
         return self.kv_buffer[:, :, indices]
 
+    def assign_flat_data(self, indices, flat_data):
+        self.kv_buffer[:, :, indices] = flat_data
+
     @debug_timing
     def transfer(self, indices, flat_data):
         # backup prepared data from device to host