chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 Packages Projects Releases Wiki Activity

Refactor prefix cache type checking (#17028)

Browse Source
This commit is contained in:
Ke Bao
2026-01-15 11:28:13 +08:00
committed by GitHub
parent 6b065298b5
commit 7f8a58fffb
12 changed files with 71 additions and 36 deletions
Download Patch File Download Diff File Expand all files Collapse all files

16
python/sglang/bench_one_batch.py
View File

@@ -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,

3
python/sglang/srt/managers/schedule_batch.py
View File

@@ -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 {}
),
)

13
python/sglang/srt/managers/schedule_policy.py
View File

@@ -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)

22
python/sglang/srt/managers/scheduler_runtime_checker_mixin.py
View File

@@ -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()

12
python/sglang/srt/mem_cache/base_prefix_cache.py
View File

@@ -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()

6
python/sglang/srt/mem_cache/chunk_cache.py
View File

@@ -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,

16
python/sglang/srt/mem_cache/common.py
View File

@@ -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

3
python/sglang/srt/mem_cache/mamba_radix_cache.py
View File

@@ -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)

3
python/sglang/srt/mem_cache/swa_radix_cache.py
View File

@@ -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 = []

3
python/sglang/srt/speculative/eagle_info_v2.py
View File

@@ -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)

3
python/sglang/srt/speculative/eagle_worker.py
View File

@@ -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)