sglang 0.5.3rc0__py3-none-any.whl → 0.5.4__py3-none-any.whl

sglang/srt/mem_cache/storage/eic/test_unit.py

@@ -0,0 +1,115 @@
+import argparse
+import os
+
+import eic
+import torch
+import yaml
+
+
+def pase_args():
+    parser = argparse.ArgumentParser(description="EIC Storage Unit Test")
+    parser.add_argument(
+        "--config",
+        "-c",
+        type=str,
+        default="/sgl-workspace/config/remote-eic.yaml",
+        help="EIC yaml config",
+    )
+    args, _ = parser.parse_known_args()
+    return args
+
+
+def init_eic_client():
+    args = pase_args()
+    config_path = os.path.abspath(args.config)
+    if not os.path.exists(config_path):
+        raise FileNotFoundError(f"Config file not found: {config_path}")
+    with open(config_path, "r") as fin:
+        config = yaml.safe_load(fin)
+
+    remote_url = config.get("remote_url", None)
+    if remote_url is None:
+        AssertionError("remote_url is None")
+    endpoint = remote_url[len("eic://") :]
+    eic_instance_id = config.get("eic_instance_id", None)
+    eic_log_dir = config.get("eic_log_dir", None)
+    eic_log_level = config.get("eic_log_level", 2)
+    eic_trans_type = config.get("eic_trans_type", 3)
+    eic_flag_file = config.get("eic_flag_file", None)
+
+    if not os.path.exists(eic_log_dir):
+        os.makedirs(eic_log_dir, exist_ok=True)
+    eic_client = eic.Client()
+    init_option = eic.InitOption()
+    init_option.log_dir = eic_log_dir
+    init_option.log_level = eic.LogLevel(eic_log_level)
+    init_option.transport_type = eic.TransportType(eic_trans_type)
+    init_option.flag_file = eic_flag_file
+    ret = eic_client.init(eic_instance_id, endpoint, init_option)
+    if ret != 0:
+        raise RuntimeError(f"EIC Client init failed with error code: {ret}")
+    return eic_client
+
+
+def test_set(eic_client):
+    test_key = ["test_key_" + str(i) for i in range(16)]
+    tensors = [
+        torch.ones([12, 6, 1, 512], dtype=torch.bfloat16, device="cpu")
+        for _ in range(16)
+    ]
+    data_keys = eic.StringVector()
+    data_vals = eic.IOBuffers()
+    for i in range(16):
+        data_keys.append(test_key[i])
+        data_vals.append(
+            tensors[i].data_ptr(), tensors[i].numel() * tensors[i].element_size(), False
+        )
+    set_opt = eic.SetOption()
+    set_opt.ttl_second = 3
+    status_code, set_outcome = eic_client.mset(data_keys, data_vals, set_opt)
+    assert (
+        status_code == eic.StatusCode.SUCCESS
+    ), f"Set failed with status code: {status_code}"
+
+
+def test_get(eic_client):
+    test_key = ["test_key_" + str(i) for i in range(16)]
+    tensors = [
+        torch.zeros([12, 6, 1, 512], dtype=torch.bfloat16, device="cpu")
+        for _ in range(16)
+    ]
+    data_keys = eic.StringVector()
+    data_vals = eic.IOBuffers()
+    for i in range(16):
+        data_keys.append(test_key[i])
+        data_vals.append(
+            tensors[i].data_ptr(), tensors[i].numel() * tensors[i].element_size(), False
+        )
+    get_opt = eic.GetOption()
+    status_code, data_vals, get_outcome = eic_client.mget(data_keys, get_opt, data_vals)
+    assert (
+        status_code == eic.StatusCode.SUCCESS
+    ), f"Get failed with status code: {status_code}"
+
+
+def test_exists(eic_client):
+    test_key = ["test_key_" + str(i) for i in range(16)]
+    data_keys = eic.StringVector()
+    for key in test_key:
+        data_keys.append(key)
+    exists_opt = eic.ExistOption()
+    status_code, exists_outcome = eic_client.mexist(data_keys, exists_opt)
+    assert (
+        status_code == eic.StatusCode.SUCCESS
+    ), f"Exists failed with status code: {status_code}"
+
+
+def main():
+    eic_client = init_eic_client()
+    test_set(eic_client)
+    test_exists(eic_client)
+    test_get(eic_client)
+
+
+if __name__ == "__main__":
+    main()

sglang/srt/mem_cache/storage/hf3fs/hf3fs_client.py

@@ -1,6 +1,5 @@
 import logging
 import os
-import threading
 from abc import ABC, abstractmethod
 from typing import List
 

sglang/srt/mem_cache/storage/hf3fs/storage_hf3fs.py

@@ -12,7 +12,12 @@ from typing import Any, List, Optional, Tuple
 
 import torch
 
-from sglang.srt.mem_cache.hicache_storage import HiCacheStorage, HiCacheStorageConfig
+from sglang.srt.mem_cache.hicache_storage import (
+    HiCacheStorage,
+    HiCacheStorageConfig,
+    HiCacheStorageExtraInfo,
+)
+from sglang.srt.mem_cache.memory_pool_host import HostKVCache
 from sglang.srt.mem_cache.storage.hf3fs.hf3fs_client import Hf3fsClient
 from sglang.srt.metrics.collector import StorageMetrics
 
@@ -178,11 +183,14 @@ class HiCacheHF3FS(HiCacheStorage):
             self.skip_backup = True
             self.rank = 0
 
+        self.is_zero_copy = False
+
         logger.info(
             f"[Rank {self.rank}] HiCacheHF3FS Client Initializing: "
             f"file_path={self.file_path}, "
             f"file_size={self.file_size / (2 ** 30):.2f} GB, "
-            f"num_pages={self.num_pages}"
+            f"num_pages={self.num_pages}, "
+            f"is_mla_model={self.is_mla_model}"
         )
 
         self.ac = AtomicCounter(self.numjobs)
@@ -323,25 +331,12 @@ class HiCacheHF3FS(HiCacheStorage):
             use_mock_client=use_mock_client,
         )
 
-    def get(
-        self,
-        key: str,
-        target_location: Optional[Any] = None,
-        target_sizes: Optional[Any] = None,
-    ) -> torch.Tensor | None:
-        return self.batch_get(
-            [key],
-            [target_location] if target_location is not None else None,
-            [target_sizes] if target_sizes is not None else None,
-        )[0]
-
     @synchronized()
-    def batch_get(
+    def _batch_get(
         self,
         keys: List[str],
-        target_locations: Optional[Any] = None,
-        target_sizes: Optional[Any] = None,
-    ) -> List[torch.Tensor | None]:
+        values: List[torch.Tensor],
+    ) -> List[bool]:
         page_indices = self.metadata_client.get_page_indices(self.rank, keys)
 
         batch_indices, file_offsets = [], []
@@ -350,15 +345,9 @@ class HiCacheHF3FS(HiCacheStorage):
                 batch_indices.append(i)
                 file_offsets.append(page_index * self.bytes_per_page)
 
-        if target_locations is not None:
-            for target_location in target_locations:
-                assert target_location.is_contiguous()
-            file_results = target_locations
-        else:
-            file_results = [
-                torch.empty(self.numel, dtype=self.dtype)
-                for _ in range(len(batch_indices))
-            ]
+        for target_location in values:
+            assert target_location.is_contiguous()
+        file_results = values
 
         start_time = time.perf_counter()
 
@@ -379,12 +368,10 @@ class HiCacheHF3FS(HiCacheStorage):
             ionum / (end_time - start_time) * self.gb_per_page
         )
 
-        results = [None] * len(keys)
-        for batch_index, file_result, read_result in zip(
-            batch_indices, file_results, read_results
-        ):
+        results = [False] * len(keys)
+        for batch_index, read_result in zip(batch_indices, read_results):
             if read_result == self.bytes_per_page:
-                results[batch_index] = file_result
+                results[batch_index] = True
             else:
                 logger.error(
                     f"[Rank {self.rank}] HiCacheHF3FS get {keys[batch_index]} failed"
@@ -392,28 +379,12 @@ class HiCacheHF3FS(HiCacheStorage):
 
         return results
 
-    def set(
-        self,
-        key: str,
-        value: Optional[Any] = None,
-        target_location: Optional[Any] = None,
-        target_sizes: Optional[Any] = None,
-    ) -> bool:
-        return self.batch_set(
-            [key],
-            [value] if value is not None else None,
-            [target_location] if target_location is not None else None,
-            [target_sizes] if target_sizes is not None else None,
-        )
-
     @synchronized()
-    def batch_set(
+    def _batch_set(
         self,
         keys: List[str],
         values: Optional[Any] = None,
-        target_locations: Optional[Any] = None,
-        target_sizes: Optional[Any] = None,
-    ) -> bool:
+    ) -> List[bool]:
         # In MLA backend, only one rank needs to backup the KV cache
         if self.skip_backup:
             return True
@@ -474,7 +445,7 @@ class HiCacheHF3FS(HiCacheStorage):
                 self.rank, written_keys_to_confirm, pages_to_release
             )
 
-        return all(results)
+        return results
 
     def delete(self, key: str) -> None:
         self.metadata_client.delete_keys(self.rank, [key])
@@ -483,22 +454,28 @@ class HiCacheHF3FS(HiCacheStorage):
         result = self.metadata_client.exists(self.rank, [key])
         return result[0] if result else False
 
-    def batch_exists(self, keys: List[str]) -> int:
+    def batch_exists(
+        self, keys: List[str], extra_info: Optional[HiCacheStorageExtraInfo] = None
+    ) -> int:
+        factor = 1
+        if self.is_zero_copy and not self.is_mla_model:
+            keys = self._get_mha_zero_copy_keys(keys)
+            factor = 2
+
         results = self.metadata_client.exists(self.rank, keys)
-        for i in range(len(keys)):
-            if not results[i]:
-                return i
 
-        return len(keys)
+        i = 0
+        while i < len(keys) and results[i]:
+            i += 1
 
-    def clear(self) -> bool:
+        return i // factor
+
+    def clear(self) -> None:
         try:
             self.metadata_client.clear(self.rank)
             logger.info(f"Cleared HiCacheHF3FS for rank {self.rank}")
-            return True
         except Exception as e:
             logger.error(f"Failed to clear HiCacheHF3FS: {e}")
-            return False
 
     def close(self) -> None:
         try:
@@ -521,3 +498,147 @@ class HiCacheHF3FS(HiCacheStorage):
         self.prefetch_bandwidth.clear()
         self.backup_bandwidth.clear()
         return storage_metrics
+
+    def register_mem_pool_host(self, mem_pool_host: HostKVCache):
+        super().register_mem_pool_host(mem_pool_host)
+        self.is_zero_copy = self.mem_pool_host.layout in [
+            "page_first",
+            "page_first_direct",
+        ]
+
+        logger.info(f"{self.is_zero_copy=}, layout={self.mem_pool_host.layout}")
+
+    def _get_mha_zero_copy_keys(self, keys: List[str]) -> List[str]:
+        _keys = []
+        for k in keys:
+            _keys.append(f"{k}-k")
+            _keys.append(f"{k}-v")
+        return _keys
+
+    def _get_mha_zero_copy_values(
+        self, values: List[torch.Tensor]
+    ) -> List[torch.Tensor]:
+        _values = []
+        for value in values:
+            _values.append(value[0])
+            _values.append(value[1])
+        return _values
+
+    def _batch_get_preprocess(self, keys, host_indices):
+        page_num = len(host_indices) // self.mem_pool_host.page_size
+        # host_indices to kv_buffer
+        flat = not self.is_zero_copy
+        values = (
+            [
+                self.mem_pool_host.get_data_page(
+                    host_indices[i * self.mem_pool_host.page_size], flat=flat
+                )
+                for i in range(page_num)
+            ]
+            if self.is_zero_copy
+            else [
+                self.mem_pool_host.get_dummy_flat_data_page() for _ in range(page_num)
+            ]
+        )
+
+        if self.is_zero_copy and not self.is_mla_model:
+            keys = self._get_mha_zero_copy_keys(keys)
+            values = self._get_mha_zero_copy_values(values)
+
+        return keys, values
+
+    def _batch_get_postprocess(self, host_indices, values, results):
+        page_num = len(host_indices) // self.mem_pool_host.page_size
+
+        if self.is_zero_copy:
+            if not self.is_mla_model:
+                results = [
+                    (results[2 * i] and results[2 * i + 1]) for i in range(page_num)
+                ]
+                results = results[:page_num]
+            return results
+
+        for i in range(page_num):
+            if not results[i]:
+                break
+            self.mem_pool_host.set_from_flat_data_page(
+                host_indices[i * self.mem_pool_host.page_size], values[i]
+            )
+
+        return results
+
+    def batch_get_v1(
+        self,
+        keys: List[str],
+        host_indices: torch.Tensor,
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> List[bool]:
+        keys, values = self._batch_get_preprocess(keys, host_indices)
+        results = self._batch_get(keys, values)
+        return self._batch_get_postprocess(host_indices, values, results)
+
+    def _batch_set_preprocess(self, keys, host_indices):
+        page_num = len(host_indices) // self.mem_pool_host.page_size
+        # host_indices to kv_buffer
+        flat = not self.is_zero_copy
+        values = [
+            self.mem_pool_host.get_data_page(
+                host_indices[i * self.mem_pool_host.page_size], flat=flat
+            )
+            for i in range(page_num)
+        ]
+
+        if self.is_zero_copy and not self.is_mla_model:
+            keys = self._get_mha_zero_copy_keys(keys)
+            values = self._get_mha_zero_copy_values(values)
+
+        return keys, values
+
+    def batch_set_v1(
+        self,
+        keys: List[str],
+        host_indices: torch.Tensor,
+        extra_info: Optional[HiCacheStorageExtraInfo] = None,
+    ) -> List[bool]:
+        len_keys = len(keys)
+        keys, values = self._batch_set_preprocess(keys, host_indices)
+        results = self._batch_set(keys, values)
+        return results
+
+    # Deprecated
+    def get(
+        self,
+        key: str,
+        target_location: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> torch.Tensor | None:
+        pass
+
+    # Deprecated
+    def batch_get(
+        self,
+        keys: List[str],
+        target_locations: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> List[torch.Tensor | None] | int:
+        pass
+
+    # Deprecated
+    def set(
+        self,
+        key: str,
+        value: Optional[Any] = None,
+        target_location: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> bool:
+        pass
+
+    # Deprecated
+    def batch_set(
+        self,
+        keys: List[str],
+        values: Optional[Any] = None,
+        target_locations: Optional[Any] = None,
+        target_sizes: Optional[Any] = None,
+    ) -> bool:
+        pass

sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py

@@ -2,14 +2,14 @@ from __future__ import annotations
 
 import logging
 import threading
-from typing import TYPE_CHECKING, List, Optional
+from typing import TYPE_CHECKING, Optional
 
 import torch
 
 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 ReqToTokenPool
-from sglang.srt.mem_cache.radix_cache import RadixCache, TreeNode
+from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
 
 try:
     from lmcache.integration.sglang.sglang_adapter import (
@@ -78,6 +78,7 @@ class LMCRadixCache(RadixCache):
         tp_size: int = 1,
         rank: int = 0,
         tp_group: Optional[torch.distributed.ProcessGroup] = None,
+        eviction_policy: str = "lru",
     ):
         super().__init__(
             req_to_token_pool=req_to_token_pool,
@@ -85,6 +86,7 @@ class LMCRadixCache(RadixCache):
             page_size=page_size,
             disable=disable,
             enable_kv_cache_events=enable_kv_cache_events,
+            eviction_policy=eviction_policy,
         )
 
         kvcache = self.token_to_kv_pool_allocator.get_kvcache()
@@ -129,7 +131,7 @@ class LMCRadixCache(RadixCache):
             with self._node_lock:
                 self._in_flight_nodes.clear()
 
-    def match_prefix(self, key: List[int], **kwargs) -> MatchResult:  # type: ignore[override]
+    def match_prefix(self, key: RadixKey, **kwargs) -> MatchResult:  # type: ignore[override]
         """Match cached prefix; if there's a tail miss, prefetch from LMCache.
 
         Reuses the base matching logic to obtain (value, last_node). If there
@@ -176,7 +178,7 @@ class LMCRadixCache(RadixCache):
         with torch.cuda.stream(self.load_stream):
             num_retrieved = self.lmcache_connector.start_load_kv(
                 LoadMetadata(
-                    token_ids=key,  # full page-aligned key
+                    token_ids=key.token_ids,  # full page-aligned key
                     slot_mapping=slot_mapping,
                     offset=value.numel() - prefix_pad,  # LMCache offset convention
                 )
@@ -215,17 +217,19 @@ class LMCRadixCache(RadixCache):
 
         return base_res
 
-    def cache_finished_req(self, req: "Req") -> None:  # type: ignore[override]
+    def cache_finished_req(self, req: "Req", is_insert: bool = True) -> None:  # type: ignore[override]
         """On request completion, insert device KV into radix and store to LMCache."""
 
-        super().cache_finished_req(req)
+        super().cache_finished_req(req, is_insert=is_insert)
+        if not is_insert:
+            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)
         ]
 
-        _, new_last_node, _, _ = self.match_prefix(token_ids)
+        _, new_last_node, _, _ = self.match_prefix(RadixKey(token_ids, req.extra_key))
         assert new_last_node is not None
 
         self.inc_lock_ref(new_last_node)
@@ -275,6 +279,8 @@ if __name__ == "__main__":
         rank=0,
         tp_group=None,
     )
-    cache.insert([1, 2, 3], torch.tensor([10, 11, 12], dtype=torch.int64))
-    cache.insert([1, 2, 3, 4], torch.tensor([10, 11, 12, 13], dtype=torch.int64))
+    cache.insert(RadixKey([1, 2, 3]), torch.tensor([10, 11, 12], dtype=torch.int64))
+    cache.insert(
+        RadixKey([1, 2, 3, 4]), torch.tensor([10, 11, 12, 13], dtype=torch.int64)
+    )
     cache.pretty_print()