From 7f8a58fffba9bc33ff2d0c2cd29ae74415620362 Mon Sep 17 00:00:00 2001 From: Ke Bao Date: Thu, 15 Jan 2026 11:28:13 +0800 Subject: [PATCH] Refactor prefix cache type checking (#17028) --- python/sglang/bench_one_batch.py | 16 +++++++++++++- python/sglang/srt/managers/schedule_batch.py | 3 +-- python/sglang/srt/managers/schedule_policy.py | 13 ++++++----- .../scheduler_runtime_checker_mixin.py | 22 +++++++++---------- .../sglang/srt/mem_cache/base_prefix_cache.py | 12 ++++++++++ python/sglang/srt/mem_cache/chunk_cache.py | 6 +++++ python/sglang/srt/mem_cache/common.py | 16 ++++++-------- .../sglang/srt/mem_cache/mamba_radix_cache.py | 3 +++ .../sglang/srt/mem_cache/swa_radix_cache.py | 3 +++ .../sglang/srt/speculative/eagle_info_v2.py | 3 +-- python/sglang/srt/speculative/eagle_worker.py | 3 +-- test/manual/test_forward_split_prefill.py | 7 +++--- 12 files changed, 71 insertions(+), 36 deletions(-) diff --git a/python/sglang/bench_one_batch.py b/python/sglang/bench_one_batch.py index b3163fbdf..826753907 100644 --- a/python/sglang/bench_one_batch.py +++ b/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, diff --git a/python/sglang/srt/managers/schedule_batch.py b/python/sglang/srt/managers/schedule_batch.py index 22e173192..5fae27b64 100644 --- a/python/sglang/srt/managers/schedule_batch.py +++ b/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 {} ), ) diff --git a/python/sglang/srt/managers/schedule_policy.py b/python/sglang/srt/managers/schedule_policy.py index 0698a59a9..fa989650e 100644 --- a/python/sglang/srt/managers/schedule_policy.py +++ b/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): - self.tree_cache.dec_lock_ref(last_node, ret) + if self.tree_cache.supports_swa() and self.tree_cache.is_tree_cache(): + self.tree_cache.dec_lock_ref(last_node, swa_uuid_for_lock) else: self.tree_cache.dec_lock_ref(last_node) diff --git a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py b/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py index 36b17556c..484a949f5 100644 --- a/python/sglang/srt/managers/scheduler_runtime_checker_mixin.py +++ b/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 ( - self.is_hybrid_ssm and isinstance(self.tree_cache, MambaRadixCache) + if ( + 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( - scheduler.tree_cache, MambaRadixCache - ): + elif scheduler.is_hybrid_ssm and scheduler.tree_cache.supports_mamba(): _, info_msg = scheduler._check_mamba_memory() else: _, info_msg = scheduler._check_radix_cache_memory() diff --git a/python/sglang/srt/mem_cache/base_prefix_cache.py b/python/sglang/srt/mem_cache/base_prefix_cache.py index 826bfa6c5..8f2560d11 100644 --- a/python/sglang/srt/mem_cache/base_prefix_cache.py +++ b/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() diff --git a/python/sglang/srt/mem_cache/chunk_cache.py b/python/sglang/srt/mem_cache/chunk_cache.py index 0d5fecf46..38dbdd6fe 100644 --- a/python/sglang/srt/mem_cache/chunk_cache.py +++ b/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, diff --git a/python/sglang/srt/mem_cache/common.py b/python/sglang/srt/mem_cache/common.py index f5b001a88..3cac71921 100644 --- a/python/sglang/srt/mem_cache/common.py +++ b/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 diff --git a/python/sglang/srt/mem_cache/mamba_radix_cache.py b/python/sglang/srt/mem_cache/mamba_radix_cache.py index 5c61299cf..598dd348f 100644 --- a/python/sglang/srt/mem_cache/mamba_radix_cache.py +++ b/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) diff --git a/python/sglang/srt/mem_cache/swa_radix_cache.py b/python/sglang/srt/mem_cache/swa_radix_cache.py index 625a8a166..1757f01ed 100644 --- a/python/sglang/srt/mem_cache/swa_radix_cache.py +++ b/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 = [] diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index a14cddb05..c943a5ec9 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/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) diff --git a/python/sglang/srt/speculative/eagle_worker.py b/python/sglang/srt/speculative/eagle_worker.py index a702df4f8..9c29b3b1a 100644 --- a/python/sglang/srt/speculative/eagle_worker.py +++ b/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) diff --git a/test/manual/test_forward_split_prefill.py b/test/manual/test_forward_split_prefill.py index cf3ecb1c1..66e3262ba 100644 --- a/test/manual/test_forward_split_prefill.py +++ b/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()