Support page size > 1 (#4356)

python/sglang/srt/mem_cache/chunk_cache.py

@@ -1,7 +1,8 @@
 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, Any, Callable, List, Tuple
 
 import torch
 
@@ -24,73 +25,40 @@ class ChunkCache(BasePrefixCache):
         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_allocator = token_to_kv_pool_allocator
-        self.entries: Dict[str, ChunkCacheEntry] = {}
-
-        self.reset()
 
     def reset(self):
-        self.entries = {}
+        pass
 
-    def match_prefix(self, rid: int, key: List[int]) -> Tuple[List[int], int]:
-        if rid not in self.entries:
-            return [], None
-
-        entry = self.entries[rid]
-        max_prefix_len = len(key)
-        return entry.value[:max_prefix_len], entry
-
-    def cache_finished_req(self, req: Req, token_ids: Optional[List[int]] = None):
-        if token_ids is None:
-            token_id_len = len(req.origin_input_ids) + len(req.output_ids) - 1
-        else:
-            token_id_len = len(token_ids)
+    def match_prefix(self, **unused_kwargs) -> Tuple[List[int], int]:
+        return [], None
 
+    def cache_finished_req(self, req: Req):
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :token_id_len
+            req.req_pool_idx, : len(req.origin_input_ids) + len(req.output_ids) - 1
         ]
         self.req_to_token_pool.free(req.req_pool_idx)
         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_id_len = len(req.fill_ids)
-
         kv_indices = self.req_to_token_pool.req_to_token[
-            req.req_pool_idx, :token_id_len
+            req.req_pool_idx, : len(req.fill_ids)
         ]
 
-        if req.rid not in self.entries:
-            self.entries[req.rid] = ChunkCacheEntry(req.rid, kv_indices)
-
-        entry = self.entries[req.rid]
-        entry.value = kv_indices
+        # `req.prefix_indices` will be used in `PrefillAdder::add_chunked_req` later
         req.prefix_indices = kv_indices
-        req.last_node = entry
 
     def insert(self):
         raise NotImplementedError()
 
-    def evict(self, num_tokens: int, evict_callback: Callable):
+    def evict(self, num_tokens: int):
         pass
 
-    def inc_lock_ref(self, node):
+    def inc_lock_ref(self, node: Any):
         return 0
 
-    def dec_lock_ref(self, node):
-        return 0
-
-    def evictable_size(self):
-        return 0
-
-    def pretty_print(self):
-        return ""
-
-    def protected_size(self):
+    def dec_lock_ref(self, node: Any):
         return 0
 
     def pretty_print(self):