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

Refactor cache init logic (#13800)

Browse Source
This commit is contained in:
Liangsheng Yin
2025-11-24 11:41:46 +08:00
committed by GitHub
parent 75222bfed9
commit b2f7b08c49
17 changed files with 307 additions and 419 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
python/sglang/srt/environ.py
View File

@@ -282,6 +282,7 @@ class Envs:
SGLANG_ENABLE_OVERLAP_PLAN_STREAM = EnvBool(False)
# VLM
SGLANG_VLM_CACHE_SIZE_MB = EnvInt(100)
SGLANG_IMAGE_MAX_PIXELS = EnvInt(16384 * 28 * 28)
SGLANG_RESIZE_RESAMPLE = EnvStr("")

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

@@ -95,12 +95,7 @@ class SchedulePolicy:
self.schedule_low_priority_values_first = schedule_low_priority_values_first
# It is used to find the matching prefix for in-batch prefix caching.
self.waiting_queue_radix_tree = RadixCache(
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
page_size=1,
disable=False,
)
self.waiting_queue_radix_tree = RadixCache.create_simulated()
def calc_priority(self, waiting_queue: List[Req]) -> bool:
if self.policy == CacheAgnosticPolicy.FCFS:

129
python/sglang/srt/managers/scheduler.py
View File

@@ -148,12 +148,9 @@ from sglang.srt.managers.scheduler_update_weights_mixin import (
)
from sglang.srt.managers.session_controller import Session
from sglang.srt.managers.utils import GenerationBatchResult, validate_input_length
from sglang.srt.mem_cache.chunk_cache import ChunkCache, SWAChunkCache
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.common import release_kv_cache
from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
from sglang.srt.mem_cache.radix_cache import RadixCache
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.multiplex.multiplexing_mixin import SchedulerMultiplexMixin
from sglang.srt.parser.reasoning_parser import ReasoningParser
@@ -419,8 +416,8 @@ class Scheduler(
# Init metrics stats
self.init_metrics(tp_rank, pp_rank, dp_rank)
# Init memory pool and cache
self.init_memory_pool_and_cache()
# Init cache using the existing memory pool
self.init_cache_with_memory_pool()
# Init running status
self.waiting_queue: List[Req] = []
@@ -693,117 +690,81 @@ class Scheduler(
reasoning_parser.detector.think_end_token, add_special_tokens=False
)[0]
def init_memory_pool_and_cache(self):
def init_cache_with_memory_pool(self):
server_args = self.server_args
self.req_to_token_pool, self.token_to_kv_pool_allocator = (
self.tp_worker.get_memory_pool()
)
params = CacheInitParams(
disable=server_args.disable_radix_cache,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
page_size=self.page_size,
is_eagle=self.spec_algorithm.is_eagle(),
tp_cache_group=(
self.attn_tp_cpu_group
if self.server_args.enable_dp_attention
else self.tp_cpu_group
),
eviction_policy=server_args.radix_eviction_policy,
enable_metrics=self.enable_metrics,
enable_kv_cache_events=self.enable_kv_cache_events,
)
if (
server_args.chunked_prefill_size is not None
and server_args.disable_radix_cache
):
if self.is_hybrid:
ChunkCacheClass = SWAChunkCache
if not self.is_hybrid:
from sglang.srt.mem_cache.chunk_cache import ChunkCache
self.tree_cache = ChunkCache(params)
else:
ChunkCacheClass = ChunkCache
self.tree_cache = ChunkCacheClass(
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
page_size=self.page_size,
)
from sglang.srt.mem_cache.chunk_cache import SWAChunkCache
self.tree_cache = SWAChunkCache(params)
else:
if envs.SGLANG_EXPERIMENTAL_CPP_RADIX_TREE.get():
# lazy import to avoid JIT overhead
from sglang.srt.mem_cache.radix_cache_cpp import RadixCacheCpp
logger.info("Using experimental C++ radix tree implementation.")
self.tree_cache = RadixCacheCpp(
disable=False,
use_hicache=self.enable_hierarchical_cache,
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
tp_cache_group=self.tp_cpu_group,
page_size=self.page_size,
hicache_ratio=server_args.hicache_ratio,
hicache_size=server_args.hicache_size,
hicache_write_policy=server_args.hicache_write_policy,
enable_metrics=self.enable_metrics,
enable_kv_cache_events=self.enable_kv_cache_events,
)
self.tree_cache = RadixCacheCpp(params=params, server_args=server_args)
elif self.enable_hierarchical_cache:
self.tree_cache = HiRadixCache(
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
tp_cache_group=(
self.attn_tp_cpu_group
if self.server_args.enable_dp_attention
else self.tp_cpu_group
),
page_size=self.page_size,
eviction_policy=server_args.radix_eviction_policy,
hicache_ratio=server_args.hicache_ratio,
hicache_size=server_args.hicache_size,
hicache_write_policy=server_args.hicache_write_policy,
hicache_io_backend=server_args.hicache_io_backend,
hicache_mem_layout=server_args.hicache_mem_layout,
enable_metrics=self.enable_metrics,
hicache_storage_backend=server_args.hicache_storage_backend,
hicache_storage_prefetch_policy=server_args.hicache_storage_prefetch_policy,
model_name=server_args.served_model_name,
storage_backend_extra_config=server_args.hicache_storage_backend_extra_config,
is_eagle=self.spec_algorithm.is_eagle(),
)
from sglang.srt.mem_cache.hiradix_cache import HiRadixCache
self.tree_cache = HiRadixCache(params=params, server_args=server_args)
self.tp_worker.register_hicache_layer_transfer_counter(
self.tree_cache.cache_controller.layer_done_counter
)
elif self.is_hybrid:
from sglang.srt.mem_cache.swa_radix_cache import SWARadixCache
self.tree_cache = SWARadixCache(
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
sliding_window_size=self.sliding_window_size,
page_size=self.page_size,
disable=server_args.disable_radix_cache,
is_eagle=self.spec_algorithm.is_eagle(),
enable_metrics=self.enable_metrics,
params=params, sliding_window_size=self.sliding_window_size
)
elif self.is_hybrid_gdn:
self.tree_cache = MambaRadixCache(
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
page_size=self.page_size,
disable=server_args.disable_radix_cache,
enable_metrics=self.enable_metrics,
)
from sglang.srt.mem_cache.mamba_radix_cache import MambaRadixCache
self.tree_cache = MambaRadixCache(params)
elif server_args.enable_lmcache:
from sglang.srt.mem_cache.storage.lmcache.lmc_radix_cache import (
LMCRadixCache,
)
self.tree_cache = LMCRadixCache(
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
page_size=self.page_size,
disable=server_args.disable_radix_cache,
enable_metrics=self.enable_metrics,
params=params,
model_config=self.model_config,
tp_size=self.tp_size,
rank=self.tp_rank,
tp_group=self.tp_group,
eviction_policy=server_args.radix_eviction_policy,
)
else:
self.tree_cache = RadixCache(
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
page_size=self.page_size,
disable=server_args.disable_radix_cache,
enable_metrics=self.enable_metrics,
enable_kv_cache_events=self.enable_kv_cache_events,
eviction_policy=server_args.radix_eviction_policy,
is_eagle=self.spec_algorithm.is_eagle(),
)
self.tree_cache = RadixCache(params)
if (
server_args.disaggregation_mode == "decode"
@@ -812,11 +773,7 @@ class Scheduler(
self.decode_offload_manager = DecodeKVCacheOffloadManager(
req_to_token_pool=self.req_to_token_pool,
token_to_kv_pool_allocator=self.token_to_kv_pool_allocator,
tp_group=(
self.attn_tp_cpu_group
if self.server_args.enable_dp_attention
else self.tp_cpu_group
),
tp_group=params.tp_cache_group,
tree_cache=self.tree_cache,
server_args=self.server_args,
)
@@ -835,7 +792,7 @@ class Scheduler(
)
)
embedding_cache_size = int(os.environ.get("SGLANG_VLM_CACHE_SIZE_MB", "100"))
embedding_cache_size = envs.SGLANG_VLM_CACHE_SIZE_MB.get()
init_mm_embedding_cache(embedding_cache_size * 1024 * 1024)
def init_disaggregation(self):

26
python/sglang/srt/mem_cache/cache_init_params.py Normal file
View File

@@ -0,0 +1,26 @@
from __future__ import annotations
import dataclasses
from typing import TYPE_CHECKING, Optional
import torch
if TYPE_CHECKING:
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
@dataclasses.dataclass
class CacheInitParams:
disable: bool
req_to_token_pool: ReqToTokenPool
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator
page_size: int
is_eagle: bool = False
tp_cache_group: Optional[torch.distributed.ProcessGroup] = None
eviction_policy: str = "lru"
disable_finished_insert: bool = False
enable_metrics: bool = False
enable_kv_cache_events: bool = False

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

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

102
python/sglang/srt/mem_cache/hiradix_cache.py
View File

@@ -1,20 +1,17 @@
from __future__ import annotations
import heapq
import json
import logging
import threading
import time
from typing import List, Optional
from typing import TYPE_CHECKING, List, Optional
import torch
from sglang.srt.managers.cache_controller import HiCacheController, PrefetchOperation
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import MatchResult
from sglang.srt.mem_cache.memory_pool import (
MHATokenToKVPool,
MLATokenToKVPool,
ReqToTokenPool,
)
from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool, MLATokenToKVPool
from sglang.srt.mem_cache.memory_pool_host import (
MHATokenToKVPoolHost,
MLATokenToKVPoolHost,
@@ -22,62 +19,49 @@ from sglang.srt.mem_cache.memory_pool_host import (
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
from sglang.srt.metrics.collector import StorageMetricsCollector
if TYPE_CHECKING:
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)
class HiRadixCache(RadixCache):
def __init__(
self,
req_to_token_pool: ReqToTokenPool,
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
tp_cache_group: torch.distributed.ProcessGroup,
page_size: int,
hicache_ratio: float,
hicache_size: int,
hicache_write_policy: str,
hicache_io_backend: str,
hicache_mem_layout: str,
enable_metrics: bool,
eviction_policy: str = "lru",
hicache_storage_backend: Optional[str] = None,
hicache_storage_prefetch_policy: Optional[str] = "best_effort",
model_name: Optional[str] = None,
storage_backend_extra_config: Optional[str] = None,
is_eagle: bool = False,
):
if hicache_io_backend == "direct":
if hicache_mem_layout == "page_first":
hicache_mem_layout = "page_first_direct"
def __init__(self, params: CacheInitParams, server_args: ServerArgs):
if server_args.hicache_io_backend == "direct":
# FIXME: move this logic into server_args parsing
if server_args.hicache_mem_layout == "page_first":
server_args.hicache_mem_layout = "page_first_direct"
logger.warning(
"Page first layout is not supported with direct IO backend, switching to page first direct layout"
)
self.kv_cache = token_to_kv_pool_allocator.get_kvcache()
self.page_size = params.page_size
self.kv_cache = params.token_to_kv_pool_allocator.get_kvcache()
if isinstance(self.kv_cache, MHATokenToKVPool):
self.token_to_kv_pool_host = MHATokenToKVPoolHost(
self.kv_cache,
hicache_ratio,
hicache_size,
page_size,
hicache_mem_layout,
server_args.hicache_ratio,
server_args.hicache_size,
self.page_size,
server_args.hicache_mem_layout,
)
elif isinstance(self.kv_cache, MLATokenToKVPool):
self.token_to_kv_pool_host = MLATokenToKVPoolHost(
self.kv_cache,
hicache_ratio,
hicache_size,
page_size,
hicache_mem_layout,
server_args.hicache_ratio,
server_args.hicache_size,
self.page_size,
server_args.hicache_mem_layout,
)
else:
raise ValueError(f"HiRadixCache only supports MHA and MLA yet")
self.tp_group = tp_cache_group
self.tp_group = params.tp_cache_group
self.tp_world_size = torch.distributed.get_world_size(group=self.tp_group)
self.enable_storage = hicache_storage_backend is not None
self.enable_storage_metrics = self.enable_storage and enable_metrics
self.enable_storage = server_args.hicache_storage_backend is not None
self.enable_storage_metrics = self.enable_storage and params.enable_metrics
(
extra_config,
@@ -85,35 +69,37 @@ class HiRadixCache(RadixCache):
prefetch_timeout_base,
prefetch_timeout_per_ki_token,
hicache_storage_pass_prefix_keys,
) = self._parse_storage_backend_extra_config(storage_backend_extra_config)
) = self._parse_storage_backend_extra_config(
server_args.hicache_storage_backend_extra_config
)
self.prefetch_threshold = prefetch_threshold
self.prefetch_timeout_base = prefetch_timeout_base
self.prefetch_timeout_per_page = (
page_size / 1024 * prefetch_timeout_per_ki_token
self.page_size / 1024 * prefetch_timeout_per_ki_token
)
self.hicache_storage_pass_prefix_keys = hicache_storage_pass_prefix_keys
# TODO: support more timeout check functions
self.is_prefetch_timeout = self._prefetch_timeout_check_linear_func
self.prefetch_stop_policy = hicache_storage_prefetch_policy
self.prefetch_stop_policy = server_args.hicache_storage_prefetch_policy
self.load_cache_event = threading.Event()
self.cache_controller = HiCacheController(
token_to_kv_pool_allocator,
params.token_to_kv_pool_allocator,
self.token_to_kv_pool_host,
page_size,
self.page_size,
self.tp_group,
load_cache_event=self.load_cache_event,
write_policy=hicache_write_policy,
io_backend=hicache_io_backend,
storage_backend=hicache_storage_backend,
write_policy=server_args.hicache_write_policy,
io_backend=server_args.hicache_io_backend,
storage_backend=server_args.hicache_storage_backend,
prefetch_threshold=self.prefetch_threshold,
model_name=model_name,
model_name=server_args.served_model_name,
storage_backend_extra_config=extra_config,
)
if self.enable_storage_metrics:
# TODO: support pp
labels = {
"storage_backend": hicache_storage_backend,
"storage_backend": server_args.hicache_storage_backend,
"tp_rank": self.cache_controller.tp_rank,
"dp_rank": self.cache_controller.dp_rank,
}
@@ -128,19 +114,11 @@ class HiRadixCache(RadixCache):
self.ongoing_backup = {}
# todo: dynamically adjust the threshold
self.write_through_threshold = (
1 if hicache_write_policy == "write_through" else 2
1 if server_args.hicache_write_policy == "write_through" else 2
)
self.load_back_threshold = 10
super().__init__(
req_to_token_pool,
token_to_kv_pool_allocator,
page_size,
disable=False,
eviction_policy=eviction_policy,
is_eagle=is_eagle,
enable_metrics=enable_metrics,
)
super().__init__(params=params)
def _parse_storage_backend_extra_config(
self, storage_backend_extra_config: Optional[str]

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

@@ -28,7 +28,6 @@ from numpy import float64
from sglang.srt.mem_cache.allocator import TokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
from sglang.srt.mem_cache.memory_pool import HybridReqToTokenPool
from sglang.srt.mem_cache.radix_cache import (
RadixKey,
_key_match_page_size1,
@@ -37,6 +36,7 @@ from sglang.srt.mem_cache.radix_cache import (
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
import logging
@@ -320,28 +320,23 @@ class LRUList:
class MambaRadixCache(BasePrefixCache):
def __init__(
self,
req_to_token_pool: HybridReqToTokenPool,
token_to_kv_pool_allocator: TokenToKVPoolAllocator,
page_size: int,
disable: bool = False,
enable_metrics: bool = False,
):
assert isinstance(token_to_kv_pool_allocator, TokenToKVPoolAllocator)
self.req_to_token_pool = req_to_token_pool
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
def __init__(self, params: CacheInitParams):
assert isinstance(params.token_to_kv_pool_allocator, TokenToKVPoolAllocator)
self.req_to_token_pool = params.req_to_token_pool
self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
assert page_size == 1, "Only support page_size=1 in mamba radix cache now."
self.page_size = page_size
self.disable = disable
assert (
params.page_size == 1
), "Only support page_size=1 in mamba radix cache now."
self.page_size = params.page_size
self.disable = params.disable
if self.token_to_kv_pool_allocator:
self.device = self.token_to_kv_pool_allocator.device
else:
self.device = torch.device("cpu")
if enable_metrics:
if params.enable_metrics:
self.init_metrics_collector()
self.key_match_fn = _key_match_page_size1

76
python/sglang/srt/mem_cache/radix_cache.py
View File

@@ -1,5 +1,6 @@
from __future__ import annotations
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.utils import convert_to_bigram_key
"""
@@ -26,7 +27,7 @@ import sys
import time
from collections import defaultdict
from functools import lru_cache, partial
from typing import TYPE_CHECKING, Iterator, List, Optional, Tuple, Union
from typing import TYPE_CHECKING, Any, Iterator, List, Optional, Tuple, Union
import torch
@@ -35,7 +36,6 @@ from sglang.srt.disaggregation.kv_events import (
BlockRemoved,
BlockStored,
)
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
from sglang.srt.mem_cache.evict_policy import (
EvictionStrategy,
@@ -46,7 +46,6 @@ from sglang.srt.mem_cache.evict_policy import (
MRUStrategy,
PriorityStrategy,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
@@ -187,28 +186,19 @@ def get_child_key(key: RadixKey, page_size: int = 1):
class RadixCache(BasePrefixCache):
def __init__(
self,
req_to_token_pool: ReqToTokenPool,
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
page_size: int,
disable: bool = False,
enable_metrics: bool = False,
enable_kv_cache_events: bool = False,
eviction_policy: str = "lru",
is_eagle: bool = False,
disable_finished_insert: bool = False,
):
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
self.disable = disable
self.enable_kv_cache_events = enable_kv_cache_events
self.kv_event_queue = []
self.is_eagle = is_eagle
self.disable_finished_insert = disable_finished_insert
def __init__(self, params: CacheInitParams):
self.disable = params.disable
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
self.enable_kv_cache_events = params.enable_kv_cache_events
self.is_eagle = params.is_eagle
self.disable_finished_insert = params.disable_finished_insert
self.eviction_policy = params.eviction_policy.lower()
if enable_metrics:
self.kv_event_queue = []
if params.enable_metrics:
self.init_metrics_collector()
if self.token_to_kv_pool_allocator:
@@ -220,27 +210,45 @@ class RadixCache(BasePrefixCache):
self.key_match_fn = _key_match_page_size1
self.get_child_key_fn = get_child_key
else:
self.key_match_fn = partial(_key_match_paged, page_size=page_size)
self.get_child_key_fn = partial(get_child_key, page_size=page_size)
self.key_match_fn = partial(_key_match_paged, page_size=self.page_size)
self.get_child_key_fn = partial(get_child_key, page_size=self.page_size)
if eviction_policy.lower() == "lru":
if self.eviction_policy == "lru":
self.eviction_strategy: EvictionStrategy = LRUStrategy()
elif eviction_policy.lower() == "lfu":
elif self.eviction_policy == "lfu":
self.eviction_strategy: EvictionStrategy = LFUStrategy()
elif eviction_policy.lower() == "fifo":
elif self.eviction_policy == "fifo":
self.eviction_strategy: EvictionStrategy = FIFOStrategy()
elif eviction_policy.lower() == "mru":
elif self.eviction_policy == "mru":
self.eviction_strategy: EvictionStrategy = MRUStrategy()
elif eviction_policy.lower() == "filo":
elif self.eviction_policy == "filo":
self.eviction_strategy: EvictionStrategy = FILOStrategy()
elif eviction_policy.lower() == "priority":
elif self.eviction_policy == "priority":
self.eviction_strategy: EvictionStrategy = PriorityStrategy()
else:
raise ValueError(
f"Unknown eviction policy: {eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority'."
f"Unknown eviction policy: {self.eviction_policy}. Supported policies: 'lru', 'lfu', 'fifo', 'mru', 'filo', 'priority'."
)
self.reset()
@classmethod
def create_simulated(
self,
disable: bool = False,
mock_allocator: Optional[Any] = None,
page_size: int = 1,
enable_kv_cache_events: bool = False,
) -> RadixCache:
"""Init a radix cache without memory pools for simulation purpose."""
params = CacheInitParams(
disable=disable,
req_to_token_pool=None,
token_to_kv_pool_allocator=mock_allocator,
page_size=page_size,
enable_kv_cache_events=enable_kv_cache_events,
)
return RadixCache(params)
##### Public API #####
def reset(self):
@@ -743,7 +751,7 @@ class RadixCache(BasePrefixCache):
if __name__ == "__main__":
tree = RadixCache(None, None, page_size=1, disable=False)
tree = RadixCache.create_simulated()
# Example token id sequences (as lists of ints)
tree.insert(RadixKey(token_ids=[1, 2, 3], extra_key=None))

102
python/sglang/srt/mem_cache/radix_cache_cpp.py
View File

@@ -6,24 +6,68 @@ from typing import TYPE_CHECKING, List, Set
import torch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import BasePrefixCache, MatchResult
from sglang.srt.mem_cache.cpp_radix_tree.radix_tree import (
IOHandle,
RadixTreeCpp,
TreeNodeCpp,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixKey
if TYPE_CHECKING:
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)
class RadixCacheCpp(BasePrefixCache):
def __init__(
self,
params: CacheInitParams,
server_args: ServerArgs,
enable_write_cancel: bool = False,
):
self.disable = params.disable
self.enable_write_cancel = enable_write_cancel
assert (
params.enable_kv_cache_events is False
), "HiRadixCache does not support kv cache events yet"
# record the nodes with ongoing write through
self.ongoing_write_through: Set[IOHandle] = set()
# record the node segments with ongoing load back
self.ongoing_load_back: Set[IOHandle] = set()
# todo: dynamically adjust the threshold
self.write_through_threshold = (
1 if server_args.hicache_write_policy == "write_through" else 2
)
self.device = self.token_to_kv_pool_allocator.device
self.token_to_kv_pool_allocator = params.token_to_kv_pool_allocator
self.req_to_token_pool = params.req_to_token_pool
self.page_size = params.page_size
self.kv_cache = self.token_to_kv_pool_allocator.get_kvcache()
self.tp_group = params.tp_cache_group
if params.enable_metrics:
self.init_metrics_collector()
if not server_args.enable_hierarchical_cache:
self.tree = RadixTreeCpp(
disabled=self.disable,
page_size=self.page_size,
host_size=None, # no host cache, this should be removed in the future
write_through_threshold=self.write_through_threshold,
)
self.cache_controller = None
return # early return if hicache is not used
raise NotImplementedError("Host cache is not supported yet")
def _merge_tensor(self, l: List[torch.Tensor]) -> torch.Tensor:
"""
Merge a list of tensors into a single tensor.
@@ -39,60 +83,6 @@ class RadixCacheCpp(BasePrefixCache):
else:
return torch.cat(l)
def __init__(
self,
disable: bool,
use_hicache: bool,
req_to_token_pool: ReqToTokenPool,
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
tp_cache_group: torch.distributed.ProcessGroup,
page_size: int,
hicache_ratio: float,
hicache_size: int,
hicache_write_policy: str,
enable_metrics: bool = False,
enable_kv_cache_events: bool = False,
hicache_oracle: bool = False,
enable_write_cancel: bool = False,
):
self.disable = disable
self.enable_write_cancel = enable_write_cancel
assert (
enable_kv_cache_events is False
), "HiRadixCache does not support kv cache events yet"
self.kv_cache = token_to_kv_pool_allocator.get_kvcache()
# record the nodes with ongoing write through
self.ongoing_write_through: Set[IOHandle] = set()
# record the node segments with ongoing load back
self.ongoing_load_back: Set[IOHandle] = set()
# todo: dynamically adjust the threshold
self.write_through_threshold = (
1 if hicache_write_policy == "write_through" else 2
)
self.device = token_to_kv_pool_allocator.device
self.token_to_kv_pool_allocator = token_to_kv_pool_allocator
self.req_to_token_pool = req_to_token_pool
self.page_size = page_size
self.tp_group = tp_cache_group
if enable_metrics:
self.init_metrics_collector()
if not use_hicache:
self.tree = RadixTreeCpp(
disabled=self.disable,
page_size=page_size,
host_size=None, # no host cache, this should be removed in the future
write_through_threshold=self.write_through_threshold,
)
self.cache_controller = None
return # early return if hicache is not used
raise NotImplementedError("Host cache is not supported yet")
def reset(self):
if self.cache_controller is not None:
# need to clear the acks before resetting the cache controller

28
python/sglang/srt/mem_cache/storage/lmcache/lmc_radix_cache.py
View File

@@ -6,9 +6,7 @@ from typing import TYPE_CHECKING, Optional
import torch
from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator
from sglang.srt.mem_cache.base_prefix_cache import MatchResult
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.radix_cache import RadixCache, RadixKey, TreeNode
try:
@@ -25,6 +23,7 @@ except ImportError as e:
if TYPE_CHECKING:
from sglang.srt.configs.model_config import ModelConfig
from sglang.srt.managers.schedule_batch import Req
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
logger = logging.getLogger(__name__)
@@ -69,27 +68,13 @@ class LMCRadixCache(RadixCache):
def __init__(
self,
req_to_token_pool: ReqToTokenPool,
token_to_kv_pool_allocator: BaseTokenToKVPoolAllocator,
page_size: int,
disable: bool = False,
enable_metrics: bool = False,
enable_kv_cache_events: bool = False,
params: CacheInitParams,
model_config: Optional["ModelConfig"] = None,
tp_size: int = 1,
rank: int = 0,
tp_group: Optional[torch.distributed.ProcessGroup] = None,
eviction_policy: str = "lru",
):
super().__init__(
req_to_token_pool=req_to_token_pool,
token_to_kv_pool_allocator=token_to_kv_pool_allocator,
page_size=page_size,
disable=disable,
enable_metrics=enable_metrics,
enable_kv_cache_events=enable_kv_cache_events,
eviction_policy=eviction_policy,
)
super().__init__(params)
kvcache = self.token_to_kv_pool_allocator.get_kvcache()
self.lmcache_connector = LMCacheLayerwiseConnector(
@@ -271,12 +256,17 @@ class LMCRadixCache(RadixCache):
if __name__ == "__main__":
cache = LMCRadixCache(
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
params = CacheInitParams(
req_to_token_pool=None,
token_to_kv_pool_allocator=None,
page_size=1,
disable=False,
enable_kv_cache_events=False,
)
cache = LMCRadixCache(
params=params,
model_config=None,
tp_size=1,
rank=0,

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

@@ -30,7 +30,7 @@ from numpy import float64
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
from sglang.srt.mem_cache.cache_init_params import CacheInitParams
from sglang.srt.mem_cache.radix_cache import (
RadixKey,
_key_match_page_size1,
@@ -329,22 +329,13 @@ class LRUList:
class SWARadixCache(BasePrefixCache):
def __init__(
self,
req_to_token_pool: ReqToTokenPool,
token_to_kv_pool_allocator: SWATokenToKVPoolAllocator,
sliding_window_size: int,
page_size: int,
disable: bool = False,
is_eagle: bool = False,
enable_metrics: bool = False,
):
assert isinstance(token_to_kv_pool_allocator, SWATokenToKVPoolAllocator)
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
self.disable = disable
self.is_eagle = is_eagle
def __init__(self, params: CacheInitParams, sliding_window_size: int):
assert isinstance(params.token_to_kv_pool_allocator, SWATokenToKVPoolAllocator)
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
self.disable = params.disable
self.is_eagle = params.is_eagle
if self.token_to_kv_pool_allocator:
self.device = self.token_to_kv_pool_allocator.device
@@ -355,15 +346,15 @@ class SWARadixCache(BasePrefixCache):
self.key_match_fn = _key_match_page_size1
self.get_child_key_fn = get_child_key
else:
self.key_match_fn = partial(_key_match_paged, page_size=page_size)
self.get_child_key_fn = partial(get_child_key, page_size=page_size)
self.key_match_fn = partial(_key_match_paged, page_size=self.page_size)
self.get_child_key_fn = partial(get_child_key, page_size=self.page_size)
if is_eagle:
if self.is_eagle:
self.key_convert_fn = convert_to_bigram_key
else:
self.key_convert_fn = lambda key: key
if enable_metrics:
if params.enable_metrics:
self.init_metrics_collector()
self.sliding_window_size = sliding_window_size

12
python/sglang/srt/utils/common.py
View File

@@ -56,6 +56,7 @@ from json import JSONDecodeError
from multiprocessing.reduction import ForkingPickler
from pathlib import Path
from typing import (
TYPE_CHECKING,
Any,
Callable,
Dict,
@@ -94,6 +95,9 @@ from typing_extensions import Literal
from sglang.srt.environ import envs
from sglang.srt.metrics.func_timer import enable_func_timer
if TYPE_CHECKING:
from sglang.srt.server_args import ServerArgs
logger = logging.getLogger(__name__)
show_time_cost = False
@@ -2754,7 +2758,7 @@ class Withable(Generic[T]):
self._value = None
def require_mlp_tp_gather(server_args):
def require_mlp_tp_gather(server_args: ServerArgs):
"""
Check if the input of MLP is obtained by all-gather rather than all-reduce. This only happens when each MLP TP group contains multiple attention DP groups.
"""
@@ -2777,7 +2781,7 @@ def require_mlp_tp_gather(server_args):
return False
def require_attn_tp_gather(server_args):
def require_attn_tp_gather(server_args: ServerArgs):
"""
Check if the input of attention is scattered.
"""
@@ -2791,11 +2795,11 @@ def require_attn_tp_gather(server_args):
return False
def require_gathered_buffer(server_args):
def require_gathered_buffer(server_args: ServerArgs):
return require_mlp_tp_gather(server_args) or require_attn_tp_gather(server_args)
def require_mlp_sync(server_args):
def require_mlp_sync(server_args: ServerArgs):
return server_args.enable_dp_attention or require_gathered_buffer(server_args)