sglang 0.4.9.post2__py3-none-any.whl → 0.4.9.post3__py3-none-any.whl

sglang/srt/mem_cache/memory_pool_host.py

@@ -71,11 +71,12 @@ class HostKVCache(abc.ABC):
         requested_bytes = self.size * self.size_per_token
         # preserve at least 10GB for other usage
         ten_gb = 10 * (1024**3)
-        if requested_bytes > host_mem.available - ten_gb:
+        available_bytes = host_mem.available - ten_gb
+        if requested_bytes > available_bytes:
             raise ValueError(
                 f"Not enough host memory available. Requesting "
                 f"{requested_bytes / 1e9:.2f} GB but only have "
-                f"{host_mem.available / 1e9:.2f} GB free. Please reduce the "
+                f"{available_bytes / 1e9:.2f} GB free. Please reduce the "
                 f"size of the hierarchical cache."
             )
         else:
@@ -98,6 +99,20 @@ class HostKVCache(abc.ABC):
     def init_kv_buffer(self):
         raise NotImplementedError()
 
+    @abc.abstractmethod
+    def get_flat_data_page(self, index) -> torch.Tensor:
+        """
+        Get a flat data page from the host memory pool.
+        """
+        raise NotImplementedError()
+
+    @abc.abstractmethod
+    def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
+        """
+        Set a flat data page to the host memory pool.
+        """
+        raise NotImplementedError()
+
     @synchronized()
     def clear(self):
         # Initialize memory states and tracking structures.
@@ -226,6 +241,19 @@ class MHATokenToKVPoolHost(HostKVCache):
             pin_memory=self.pin_memory,
         )
 
+    # todo, page first memory layout
+    def get_flat_data_page(self, index) -> torch.Tensor:
+        return self.kv_buffer[:, :, index : index + self.page_size, :, :].flatten()
+
+    def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
+        self.kv_buffer[:, :, index : index + self.page_size, :, :] = data_page.reshape(
+            2,
+            self.layer_num,
+            self.page_size,
+            self.head_num,
+            self.head_dim,
+        )
+
     @property
     def k_buffer(self):
         return self.kv_buffer[0]
@@ -275,3 +303,14 @@ class MLATokenToKVPoolHost(HostKVCache):
             device=self.device,
             pin_memory=self.pin_memory,
         )
+
+    def get_flat_data_page(self, index) -> torch.Tensor:
+        return self.kv_buffer[:, index : index + self.page_size, :, :].flatten()
+
+    def set_from_flat_data_page(self, index: int, data_page: torch.Tensor) -> None:
+        self.kv_buffer[:, index : index + self.page_size, :, :] = data_page.reshape(
+            self.layer_num,
+            self.page_size,
+            1,
+            self.kv_lora_rank + self.qk_rope_head_dim,
+        )

sglang/srt/mem_cache/radix_cache.py

@@ -55,8 +55,13 @@ class TreeNode:
         self.hit_count = 0
         # indicating the node is loading KV cache from host
         self.loading = False
+        # indicating the node is locked to protect from eviction
+        # incremented when the node is referenced by a storage operation
+        self.host_ref_counter = 0
         # store the host indices of KV cache
         self.host_value: Optional[torch.Tensor] = None
+        # store hash values of each pages
+        self.hash_value: Optional[List[str]] = None
 
         self.id = TreeNode.counter if id is None else id
         TreeNode.counter += 1
@@ -69,6 +74,27 @@ class TreeNode:
     def backuped(self):
         return self.host_value is not None
 
+    @property
+    def backuped_storage(self):
+        return self.hash_value is not None and len(self.hash_value) > 0
+
+    def protect_host(self):
+        """Protect the host value from eviction."""
+        self.host_ref_counter += 1
+
+    def release_host(self):
+        """Release the host value, allowing it to be evicted."""
+        if self.host_ref_counter > 0:
+            self.host_ref_counter -= 1
+        else:
+            raise RuntimeError("Host reference counter is already zero.")
+
+    def get_last_hash_value(self) -> Optional[str]:
+        """Returns the hash value of the last page in this node."""
+        if self.hash_value is None or len(self.hash_value) == 0:
+            return None
+        return self.hash_value[-1]
+
     def __lt__(self, other: "TreeNode"):
         return self.last_access_time < other.last_access_time