Refactor prefix cache type checking (#17028)

python/sglang/bench_one_batch.py

@@ -351,10 +351,24 @@ def prepare_synthetic_inputs_for_latency_test(
     return reqs
 
 
+class TreeCacheNamespace(SimpleNamespace):
+    def supports_swa(self) -> bool:
+        return False
+
+    def supports_mamba(self) -> bool:
+        return False
+
+    def is_chunk_cache(self) -> bool:
+        return False
+
+    def is_tree_cache(self) -> bool:
+        return not self.is_chunk_cache()
+
+
 @torch.no_grad
 def extend(reqs, model_runner):
     # Create dummy tree_cache for benchmarks (no prefix caching, just allocation)
-    dummy_tree_cache = SimpleNamespace(
+    dummy_tree_cache = TreeCacheNamespace(
         page_size=model_runner.server_args.page_size,
         device=model_runner.device,
         token_to_kv_pool_allocator=model_runner.token_to_kv_pool_allocator,

python/sglang/srt/managers/schedule_batch.py

@@ -66,7 +66,6 @@ from sglang.srt.mem_cache.common import (
     evict_from_tree_cache,
     release_kv_cache,
 )
-from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
 from sglang.srt.mem_cache.radix_cache import RadixKey
 from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
@@ -868,7 +867,7 @@ class Req:
                 key=RadixKey(token_ids=token_ids, extra_key=self.extra_key),
                 **(
                     {"req": self, "cow_mamba": True}
-                    if isinstance(tree_cache, MambaRadixCache)
+                    if tree_cache.supports_mamba()
                     else {}
                 ),
             )

python/sglang/srt/managers/schedule_policy.py

@@ -35,10 +35,8 @@ import torch
 from sglang.srt.layers.attention.nsa.utils import is_nsa_prefill_cp_in_seq_split
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache
-from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
 from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
 from sglang.srt.mem_cache.swa_memory_pool import SWATokenToKVPoolAllocator
-from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
 from sglang.srt.server_args import ServerArgs
 
 if TYPE_CHECKING:
@@ -407,7 +405,7 @@ class PrefillAdder:
         self.is_hybrid_swa = isinstance(
             self.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator
         )
-        self.is_hybrid_ssm_cache = isinstance(self.tree_cache, MambaRadixCache)
+        self.is_hybrid_ssm_cache = self.tree_cache.supports_mamba()
 
         self.priority_scheduling_preemption_threshold = (
             priority_scheduling_preemption_threshold
@@ -523,11 +521,14 @@ class PrefillAdder:
     @contextmanager
     def _lock_node(self, last_node: TreeNode):
         try:
-            ret = self.tree_cache.inc_lock_ref(last_node)
+            if self.tree_cache.supports_swa() and self.tree_cache.is_tree_cache():
+                swa_uuid_for_lock = self.tree_cache.inc_lock_ref(last_node)
+            else:
+                self.tree_cache.inc_lock_ref(last_node)
             yield None
         finally:
-            if isinstance(self.tree_cache, SWARadixCache):
+            if self.tree_cache.supports_swa() and self.tree_cache.is_tree_cache():
-                self.tree_cache.dec_lock_ref(last_node, ret)
+                self.tree_cache.dec_lock_ref(last_node, swa_uuid_for_lock)
             else:
                 self.tree_cache.dec_lock_ref(last_node)
 

python/sglang/srt/managers/scheduler_runtime_checker_mixin.py

@@ -8,8 +8,6 @@ from typing import TYPE_CHECKING
 from sglang.srt.disaggregation.utils import DisaggregationMode
 from sglang.srt.environ import envs
 from sglang.srt.managers.schedule_batch import ScheduleBatch
-from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
-from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
 from sglang.srt.utils.common import ceil_align, raise_error_or_warn
 from sglang.srt.utils.request_logger import disable_request_logging
 from sglang.srt.utils.watchdog import WatchdogRaw
@@ -29,14 +27,16 @@ class SchedulerRuntimeCheckerMixin:
         return num_used, token_usage, available_size, evictable_size
 
     def _get_mamba_token_info(self: Scheduler):
-        is_radix_tree = isinstance(self.tree_cache, MambaRadixCache)
+        is_mamba_radix_cache = (
+            self.tree_cache.supports_mamba() and self.tree_cache.is_tree_cache()
+        )
         full_available_size = self.token_to_kv_pool_allocator.available_size()
         full_evictable_size = (
-            self.tree_cache.full_evictable_size() if is_radix_tree else 0
+            self.tree_cache.full_evictable_size() if is_mamba_radix_cache else 0
         )
         mamba_available_size = self.req_to_token_pool.mamba_pool.available_size()
         mamba_evictable_size = (
-            self.tree_cache.mamba_evictable_size() if is_radix_tree else 0
+            self.tree_cache.mamba_evictable_size() if is_mamba_radix_cache else 0
         )
         full_num_used = self.token_to_kv_pool_allocator.size - (
             full_available_size + full_evictable_size
@@ -234,7 +234,7 @@ class SchedulerRuntimeCheckerMixin:
     def check_memory(self: Scheduler):
         if self.is_hybrid_swa:
             memory_leak, token_msg = self._check_hybrid_memory()
-        elif self.is_hybrid_ssm and isinstance(self.tree_cache, MambaRadixCache):
+        elif self.is_hybrid_ssm and self.tree_cache.supports_mamba():
             memory_leak, token_msg = self._check_mamba_memory()
         else:
             memory_leak, token_msg = self._check_radix_cache_memory()
@@ -307,8 +307,10 @@ class SchedulerRuntimeCheckerMixin:
         self._publish_kv_events()
 
     def check_tree_cache(self: Scheduler):
-        if (self.is_hybrid_swa and isinstance(self.tree_cache, SWARadixCache)) or (
+        if (
-            self.is_hybrid_ssm and isinstance(self.tree_cache, MambaRadixCache)
+            self.tree_cache.is_tree_cache()
+            and (self.is_hybrid_swa and self.tree_cache.supports_swa())
+            or (self.is_hybrid_ssm and self.tree_cache.supports_mamba())
         ):
             self.tree_cache.sanity_check()
 
@@ -341,9 +343,7 @@ def create_scheduler_watchdog(
             return ""
         if scheduler.is_hybrid_swa:
             _, info_msg = scheduler._check_hybrid_memory()
-        elif scheduler.is_hybrid_ssm and isinstance(
+        elif scheduler.is_hybrid_ssm and scheduler.tree_cache.supports_mamba():
-            scheduler.tree_cache, MambaRadixCache
-        ):
             _, info_msg = scheduler._check_mamba_memory()
         else:
             _, info_msg = scheduler._check_radix_cache_memory()

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -148,3 +148,15 @@ class BasePrefixCache(ABC, PrefixCacheTrait):
 
     def take_events(self):
         return []
+
+    def supports_swa(self) -> bool:
+        return False
+
+    def supports_mamba(self) -> bool:
+        return False
+
+    def is_chunk_cache(self) -> bool:
+        return False
+
+    def is_tree_cache(self) -> bool:
+        return not self.is_chunk_cache()

python/sglang/srt/mem_cache/chunk_cache.py

@@ -30,6 +30,9 @@ class ChunkCache(BasePrefixCache):
 
         self.protected_size_ = 0
 
+    def is_chunk_cache(self) -> bool:
+        return True
+
     # NOTE (csy): this is to determine if a cache has prefix matching feature.
     # Chunk cache always return True to indicate no prefix matching.
     # TODO (csy): Using a prefix cache trait to replace this
@@ -106,6 +109,9 @@ class SWAChunkCache(ChunkCache):
 
         self.chunked_prefill_size = params.chunked_prefill_size
 
+    def supports_swa(self) -> bool:
+        return True
+
     def evict_swa(
         self,
         req: Req,

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(
+        assert (
-            tree_cache, MambaRadixCache
+            tree_cache.supports_mamba()
         ), "Only MambaRadixCache can handle abort with prefix cache hit before alloc"
         return
 

python/sglang/srt/mem_cache/mamba_radix_cache.py

@@ -397,6 +397,9 @@ class MambaRadixCache(BasePrefixCache):
 
     ##### Public API #####
 
+    def supports_mamba(self) -> bool:
+        return True
+
     def reset(self) -> None:
         self.root_node = TreeNode()
         self.root_node.key = RadixKey([], None)

python/sglang/srt/mem_cache/swa_radix_cache.py

@@ -362,6 +362,9 @@ class SWARadixCache(BasePrefixCache):
 
     ##### Public API #####
 
+    def supports_swa(self) -> bool:
+        return True
+
     def reset(self) -> None:
         self.root_node = TreeNode()
         self.root_node.key = []

python/sglang/srt/speculative/eagle_info_v2.py

@@ -10,7 +10,6 @@ import triton.language as tl
 
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.managers.schedule_batch import ModelWorkerBatch, ScheduleBatch
-from sglang.srt.mem_cache.chunk_cache import SWAChunkCache
 from sglang.srt.mem_cache.common import (
     alloc_paged_token_slots_extend,
     alloc_token_slots,
@@ -80,7 +79,7 @@ def assign_draft_cache_locs_page_size_1(
 @dataclass
 class EagleDraftInputV2Mixin:
     def prepare_for_decode(self: EagleDraftInput, batch: ScheduleBatch):
-        if isinstance(batch.tree_cache, SWAChunkCache):
+        if batch.tree_cache.supports_swa() and batch.tree_cache.is_chunk_cache():
             for req in batch.reqs:
                 batch.tree_cache.evict_swa(req, req.seqlen - 1)
 

python/sglang/srt/speculative/eagle_worker.py

@@ -19,7 +19,6 @@ from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput
 from sglang.srt.managers.schedule_batch import ScheduleBatch
 from sglang.srt.managers.scheduler import GenerationBatchResult
 from sglang.srt.managers.tp_worker import TpModelWorker
-from sglang.srt.mem_cache.chunk_cache import SWAChunkCache
 from sglang.srt.mem_cache.common import (
     alloc_paged_token_slots_extend,
     alloc_token_slots,
@@ -375,7 +374,7 @@ class EAGLEWorker(TpModelWorker):
         )
 
     def _draft_preprocess_decode(self, batch: ScheduleBatch):
-        if isinstance(batch.tree_cache, SWAChunkCache):
+        if batch.tree_cache.supports_swa() and batch.tree_cache.is_chunk_cache():
             for req in batch.reqs:
                 batch.tree_cache.evict_swa(req, req.seqlen - 1)
 

test/manual/test_forward_split_prefill.py

@@ -8,11 +8,11 @@ python3 test_forward_split_prefill.py
 """
 
 import unittest
-from types import SimpleNamespace
 
 import numpy as np
 import torch
 
+from sglang.bench_one_batch import TreeCacheNamespace
 from sglang.srt.configs.model_config import ModelConfig
 from sglang.srt.managers.schedule_batch import Req, ScheduleBatch
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch
@@ -61,6 +61,8 @@ class TestForwardSplitPrefill(CustomTestCase):
             pp_size=1,
             nccl_port=cls.port_args.nccl_port,
             server_args=cls.server_args,
+            moe_ep_rank=0,
+            moe_ep_size=1,
         )
 
         cls.tokenizer = get_tokenizer(
@@ -97,7 +99,7 @@ class TestForwardSplitPrefill(CustomTestCase):
             reqs.append(req)
 
         # Create dummy tree_cache for tests (no prefix caching, just allocation)
-        dummy_tree_cache = SimpleNamespace(
+        dummy_tree_cache = TreeCacheNamespace(
             page_size=1,
             device=self.model_runner.device,
             token_to_kv_pool_allocator=self.model_runner.token_to_kv_pool_allocator,
@@ -111,7 +113,6 @@ class TestForwardSplitPrefill(CustomTestCase):
             model_config=self.model_config,
             enable_overlap=False,
             spec_algorithm=SpeculativeAlgorithm.NONE,
-            enable_custom_logit_processor=False,
         )
         if is_split_prefill:
             batch.prepare_for_split_prefill()