sglang 0.4.7__py3-none-any.whl → 0.4.8__py3-none-any.whl

sglang/srt/mem_cache/chunk_cache.py

@@ -2,28 +2,23 @@ from __future__ import annotations
 
 """Cache for chunked prefill, used when RadixCache is disabled."""
 
-from typing import TYPE_CHECKING, Any, Callable, List, Tuple
+from typing import TYPE_CHECKING, Any
 
 import torch
 
-from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
-from sglang.srt.mem_cache.memory_pool import ReqToTokenPool, TokenToKVPoolAllocator
+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
 
 
-class ChunkCacheEntry:
-    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_allocator: TokenToKVPoolAllocator,
+        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
         page_size: int,
     ):
         self.req_to_token_pool = req_to_token_pool
@@ -33,8 +28,12 @@ class ChunkCache(BasePrefixCache):
     def reset(self):
         pass
 
-    def match_prefix(self, **unused_kwargs) -> Tuple[List[int], int]:
-        return [], None
+    def match_prefix(self, **unused_kwargs) -> MatchResult:
+        return MatchResult(
+            device_indices=torch.empty((0,), dtype=torch.int64),
+            last_device_node=None,
+            last_host_node=None,
+        )
 
     def cache_finished_req(self, req: Req):
         kv_indices = self.req_to_token_pool.req_to_token[
@@ -53,9 +52,6 @@ class ChunkCache(BasePrefixCache):
         # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
         req.prefix_indices = kv_indices
 
-    def insert(self):
-        raise NotImplementedError()
-
     def evict(self, num_tokens: int):
         pass
 

sglang/srt/mem_cache/hiradix_cache.py

@@ -7,13 +7,16 @@ from typing import List, Optional
 import torch
 
 from sglang.srt.managers.cache_controller import HiCacheController
+from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
+from sglang.srt.mem_cache.base_prefix_cache import MatchResult
 from sglang.srt.mem_cache.memory_pool import (
     MHATokenToKVPool,
-    MHATokenToKVPoolHost,
     MLATokenToKVPool,
-    MLATokenToKVPoolHost,
     ReqToTokenPool,
-    TokenToKVPoolAllocator,
+)
+from sglang.srt.mem_cache.memory_pool_host import (
+    MHATokenToKVPoolHost,
+    MLATokenToKVPoolHost,
 )
 from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode
 
@@ -25,7 +28,7 @@ class HiRadixCache(RadixCache):
     def __init__(
         self,
         req_to_token_pool: ReqToTokenPool,
-        token_to_kv_pool_allocator: TokenToKVPoolAllocator,
+        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
         tp_cache_group: torch.distributed.ProcessGroup,
         page_size: int,
         hicache_ratio: float,
@@ -281,39 +284,44 @@ class HiRadixCache(RadixCache):
     def init_load_back(
         self,
         last_node: TreeNode,
-        prefix_indices: torch.Tensor,
+        host_hit_length: int,
         mem_quota: Optional[int] = None,
     ):
-        assert (
-            len(prefix_indices) == 0 or prefix_indices.is_cuda
-        ), "indices of device kV caches should be on GPU"
+        _ = host_hit_length  # unused, but kept for compatibility
         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}"
                 )
+                return loading_values, last_node
 
             while last_node.evicted:
                 last_node = last_node.parent
 
-        return last_node, prefix_indices
+        return (
+            torch.empty((0,), dtype=torch.int64, device=self.device),
+            last_node,
+        )
 
-    def ready_to_load_cache(self):
+    def ready_to_load_host_cache(self):
+        producer_index = self.cache_controller.layer_done_counter.next_producer()
         self.load_cache_event.set()
+        return producer_index
+
+    def check_hicache_events(self):
+        self.writing_check()
+        self.loading_check()
 
-    def match_prefix(self, key: List[int], include_evicted=False, **kwargs):
+    def match_prefix(self, key: List[int], **kwargs):
         empty_value = torch.empty((0,), dtype=torch.int64, device=self.device)
         if self.disable or len(key) == 0:
-            if include_evicted:
-                return empty_value, self.root_node, self.root_node
-            else:
-                return empty_value, self.root_node
+            return MatchResult(
+                device_indices=empty_value,
+                last_device_node=self.root_node,
+                last_host_node=self.root_node,
+                host_hit_length=0,
+            )
 
         if self.page_size != 1:
             page_aligned_len = len(key) // self.page_size * self.page_size
@@ -325,14 +333,18 @@ class HiRadixCache(RadixCache):
         else:
             value = empty_value
 
-        last_node_global = last_node
+        host_hit_length = 0
+        last_host_node = last_node
         while last_node.evicted:
+            host_hit_length += len(last_node.host_value)
             last_node = last_node.parent
 
-        if include_evicted:
-            return value, last_node, last_node_global
-        else:
-            return value, last_node
+        return MatchResult(
+            device_indices=value,
+            last_device_node=last_node,
+            last_host_node=last_host_node,
+            host_hit_length=host_hit_length,
+        )
 
     def _match_prefix_helper(self, node: TreeNode, key: List):
         node.last_access_time = time.monotonic()
@@ -370,6 +382,7 @@ class HiRadixCache(RadixCache):
         new_node.lock_ref = child.lock_ref
         new_node.key = child.key[:split_len]
         new_node.loading = child.loading
+        new_node.hit_count = child.hit_count
 
         # split value and host value if exists
         if child.evicted: