Improve type annotation (#1029)

python/sglang/srt/mem_cache/chunk_cache.py

@@ -1,8 +1,11 @@
+from __future__ import annotations
+
 """Cache for chunked prefill, used when RadixCache is disabled."""
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Callable
 
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
+from sglang.srt.mem_cache.memory_pool import BaseTokenToKVPool, ReqToTokenPool
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
@@ -15,7 +18,9 @@ class ChunkCacheEntry:
 
 
 class ChunkCache(BasePrefixCache):
-    def __init__(self, req_to_token_pool, token_to_kv_pool):
+    def __init__(
+        self, req_to_token_pool: ReqToTokenPool, token_to_kv_pool: BaseTokenToKVPool
+    ):
         self.disable = True
         self.req_to_token_pool = req_to_token_pool
         self.token_to_kv_pool = token_to_kv_pool
@@ -32,7 +37,7 @@ class ChunkCache(BasePrefixCache):
         entry = self.entries[rid]
         return entry.value, entry
 
-    def cache_finished_req(self, req: "Req", token_ids=None):
+    def cache_finished_req(self, req: Req, token_ids: Optional[List[int]] = None):
         if token_ids is None:
             token_ids = (req.origin_input_ids + req.output_ids)[:-1]
 
@@ -45,7 +50,7 @@ class ChunkCache(BasePrefixCache):
         if req.rid in self.entries:
             del self.entries[req.rid]
 
-    def cache_unfinished_req(self, req: "Req", token_ids=None):
+    def cache_unfinished_req(self, req: Req, token_ids: Optional[List[int]] = None):
         if token_ids is None:
             token_ids = req.fill_ids
 
@@ -64,7 +69,7 @@ class ChunkCache(BasePrefixCache):
     def insert(self):
         raise NotImplementedError
 
-    def evict(self, num_tokens, evict_callback):
+    def evict(self, num_tokens: int, evict_callback: Callable):
         pass
 
     def inc_lock_ref(self, node):