Refactor cache init logic (#13800)

python/sglang/srt/mem_cache/chunk_cache.py

@@ -6,27 +6,19 @@ from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 
-from sglang.srt.mem_cache.allocator import (
-    BaseTokenToKVPoolAllocator,
-    SWATokenToKVPoolAllocator,
-)
+from sglang.srt.mem_cache.allocator import SWATokenToKVPoolAllocator
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
-from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 
 if TYPE_CHECKING:
     from sglang.srt.managers.schedule_batch import Req
+    from sglang.srt.mem_cache.cache_init_params import CacheInitParams
 
 
 class ChunkCache(BasePrefixCache):
-    def __init__(
-        self,
-        req_to_token_pool: ReqToTokenPool,
-        token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
-        page_size: int,
-    ):
-        self.req_to_token_pool = req_to_token_pool
-        self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
-        self.page_size = page_size
+    def __init__(self, params: CacheInitParams):
+        self.req_to_token_pool = params.req_to_token_pool
+        self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
+        self.page_size = params.page_size
         if self.token_to_kv_pool_allocator:
             self.device = self.token_to_kv_pool_allocator.device
         else:
@@ -89,14 +81,9 @@ class ChunkCache(BasePrefixCache):
 class SWAChunkCache(ChunkCache):
     """ChunkCache with support for hybrid KV cache operations."""
 
-    def __init__(
-        self,
-        req_to_token_pool: ReqToTokenPool,
-        token_to_kv_pool_allocator: SWATokenToKVPoolAllocator,
-        page_size: int,
-    ):
-        super().__init__(req_to_token_pool, token_to_kv_pool_allocator, page_size)
-        assert isinstance(token_to_kv_pool_allocator, SWATokenToKVPoolAllocator)
+    def __init__(self, params: CacheInitParams):
+        assert isinstance(params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator)
+        super().__init__(params)
 
     def evict_swa(
         self,