[Feature] KVCache layout: compact laytout

Signed-off-by: Haiyang Shi <[email protected]>

Author: Haiyang Shi

python/aibrix_kvcache/aibrix_kvcache/__init__.py

@@ -13,7 +13,6 @@
 # limitations under the License.
 
 from .cache_handle import KVCacheHandle, MemoryRegionKVCacheHandle
-from .cache_hashable import TokenCacheKey
 from .cache_manager import (
     BaseKVCacheManager,
     GroupAwareKVCacheManager,
@@ -26,7 +25,6 @@
 
 __all__ = [
     "KVCacheHandle",
-    "TokenCacheKey",
     "MemoryRegionKVCacheHandle",
     "BaseKVCacheManager",
     "GroupAwareKVCacheManager",

python/aibrix_kvcache/aibrix_kvcache/cache_hashable.py

@@ -13,13 +13,9 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import Iterable
+from typing import Tuple
 
-import numpy as np
-from farmhash import FarmHash64
-
-from .memory import MemoryRegion
-from .utils import np_array_concat
+from .common import CachedPyObjectBase
 
 
 class KVCacheHashable(ABC):
@@ -40,119 +36,34 @@ def __len__(self) -> int:
         raise NotImplementedError
 
 
-class BaseKVCacheHashable(KVCacheHashable):
+class BaseKVCacheHashable(KVCacheHashable, CachedPyObjectBase):
     """
     Base class for a hashable object that uses all tokens to compute the hash.
     """
 
-    @abstractmethod
-    def all_tokens_memoryview(self) -> memoryview:
-        """Memoryview of the PACKED bytes representation of all tokens."""
-        raise NotImplementedError
+    def __init__(self, prefix: Tuple[int, ...] | None, tokens: Tuple[int, ...]):
+        self.prefix = prefix or tuple()
+        self.tokens = tokens
 
     def __hash__(self) -> int:
-        return FarmHash64(self.all_tokens_memoryview())
+        return hash((self.prefix, self.tokens))
 
     def __eq__(self, other) -> bool:
         if not isinstance(other, BaseKVCacheHashable):
             return False
-        return self.all_tokens_memoryview() == other.all_tokens_memoryview()
+        return (self.prefix, self.tokens) == (other.prefix, other.tokens)
 
 
 class TokenCacheKey(BaseKVCacheHashable):
     """
-    A cache key that use numpy ndarray's bytes representation to compute
-    the hash.
+    A cache key that compounds prefix and tokens.
     Args:
         prefix (np.ndarray | None): The prefix tokens of the kv tensors.
         tokens (np.ndarray): The tokens of the kv tensors.
     """
 
-    def __init__(self, *, tokens: np.ndarray, prefix: np.ndarray | None = None):
-        self._storage = np_array_concat(prefix, tokens)
-
-        self.prefix = prefix
-        self.tokens = tokens
-
-    def all_tokens(self) -> np.ndarray:
-        return self._storage
-
-    def all_tokens_memoryview(self) -> memoryview:
-        return memoryview(self._storage)  # type: ignore
-
-    def __len__(self) -> int:
-        return len(self._storage)
-
-    def shift(self, shift: int) -> "TokenCacheKey":
-        """
-        Shifts the split of prefix and tokens.
-        Args:
-            shift (int): The number of tokens to shift.
-        Returns:
-            TokenCacheKey: The cache key with shifted tokens.
-        """
-        orig_prefix_len = len(self.prefix) if self.prefix is not None else 0
-        new_prefix_len = min(max(0, orig_prefix_len + shift), len(self))
-        new_prefix = self._storage[:new_prefix_len]
-        new_tokens = self._storage[new_prefix_len:]
-        return TokenCacheKey(prefix=new_prefix, tokens=new_tokens)
-
-    def shrink(self, n: int) -> "TokenCacheKey":
-        """
-        Shrinks n from tokens.
-        Args:
-            n (int): The number of tokens to shrink.
-        Returns:
-            TokenCacheKey: The new cache key.
-        """
-        new_tokens = self.tokens[:-n]
-        return TokenCacheKey(prefix=self.prefix, tokens=new_tokens)
-
-    def batched(self, batch_size: int) -> Iterable["TokenCacheKey"]:
-        """
-        Batches the tokens into a list of TokenCacheKey.
-        Args:
-            batch_size (int): The batch size.
-        Returns:
-            Iterable[TokenCacheKey]: The batched tokens.
-        """
-        prefix_len = len(self.prefix) if self.prefix is not None else 0
-        num_batches = len(self.tokens) // batch_size
-        for _ in range(num_batches):
-            yield (
-                TokenCacheKey(
-                    prefix=self._storage[:prefix_len],
-                    tokens=self._storage[prefix_len : prefix_len + batch_size],
-                )
-            )
-            prefix_len += batch_size
-
-
-class MemoryRegionCacheEntry(BaseKVCacheHashable):
-    """
-    A cache entry that stores a memory region and uses MR's all tokens
-    to compute the hash.
-    """
-
-    def __init__(self, mr: MemoryRegion):
-        # Take the ownership of the memory region
-        self._mr: MemoryRegion | None = mr
-        assert mr.is_sealed, "Memory region must be sealed"
-
-        self.ref_up = self._mr.ref_up
-        self.ref_down = self._mr.ref_down
-
-    def all_tokens_memoryview(self) -> memoryview:
-        assert self._mr is not None
-        return memoryview(np_array_concat(*self._mr.unpack_tokens()))  # type: ignore
+    def __init__(self, prefix: Tuple[int, ...] | None, tokens: Tuple[int, ...]):
+        super().__init__(prefix, tokens)
 
     def __len__(self) -> int:
-        assert self._mr is not None
-        return self._mr.length
-
-    def cache_key(self) -> TokenCacheKey:
-        assert self._mr is not None
-        prefix, tokens = self._mr.unpack_tokens()
-        prefix = prefix.copy() if prefix is not None else None
-        tokens = tokens.copy()
-        return TokenCacheKey(prefix=prefix, tokens=tokens)
+        return len(self.prefix) + len(self.tokens)