Refactor prefix cache type checking (#17028)

python/sglang/srt/mem_cache/common.py

@@ -8,8 +8,6 @@ import triton
 import triton.language as tl
 
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
-from sglang.srt.mem_cache.chunk_cache import ChunkCache, SWAChunkCache
-from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool, ReqToTokenPool
 from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
 from sglang.srt.server_args import get_global_server_args
@@ -232,7 +230,7 @@ def evict_from_tree_cache(tree_cache: BasePrefixCache | None, num_tokens: int):
     if tree_cache is None:
         return
 
-    if isinstance(tree_cache, (SWAChunkCache, ChunkCache)):
+    if tree_cache.is_chunk_cache():
         return
 
     allocator = tree_cache.token_to_kv_pool_allocator
@@ -306,12 +304,12 @@ def alloc_req_slots(
         mamba_available_size = req_to_token_pool.mamba_pool.available_size()
         factor = (
             MAMBA_STATE_PER_REQ_PREFIX_CACHE
-            if isinstance(tree_cache, MambaRadixCache)
+            if tree_cache.supports_mamba()
             else MAMBA_STATE_PER_REQ_NO_CACHE
         )
         mamba_state_needed = num_reqs * factor
         if mamba_available_size < mamba_state_needed:
-            if tree_cache is not None and isinstance(tree_cache, MambaRadixCache):
+            if tree_cache is not None and tree_cache.supports_mamba():
                 mamba_num = max(0, mamba_state_needed - mamba_available_size)
                 tree_cache.evict_mamba(mamba_num)
         req_pool_indices = req_to_token_pool.alloc(num_reqs, reqs)
@@ -340,7 +338,7 @@ def alloc_for_extend(
         req_pool_indices: request pool indices as list
     """
     # free out-of-window swa tokens
-    if isinstance(batch.tree_cache, SWAChunkCache):
+    if batch.tree_cache.supports_swa() and batch.tree_cache.is_chunk_cache():
         for req, pre_len in zip(batch.reqs, batch.prefix_lens):
             if batch.enable_overlap:
                 # In chunked prefill case, when the second extend batch is scheduling, the first extend batch is still running, so we cannot evict swa tokens
@@ -442,7 +440,7 @@ def alloc_for_decode(batch: ScheduleBatch, token_per_req: int) -> torch.Tensor:
     Returns:
         out_cache_loc: allocated cache locations
     """
-    if isinstance(batch.tree_cache, SWAChunkCache):
+    if batch.tree_cache.supports_swa() and batch.tree_cache.is_chunk_cache():
         for req in batch.reqs:
             # We set evict_swa condition here with two reasons:
             # 1. In overlap scheduler, we cannot evict swa when req.decode_batch_idx == 0 since the prev extend batch is still running.
@@ -487,8 +485,8 @@ def release_kv_cache(req: Req, tree_cache: BasePrefixCache, is_insert: bool = Tr
 
     # MambaRadixCache may alloc mamba state before alloc KV cache
     if req.req_pool_idx is None:
-        assert isinstance(
-            tree_cache, MambaRadixCache
+        assert (
+            tree_cache.supports_mamba()
         ), "Only MambaRadixCache can handle abort with prefix cache hit before alloc"
         return