[Refactor] Let eviction policies take construction parameters (#37795)

Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

python/sglang/srt/mem_cache/cache_init_params.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 import dataclasses
-from typing import TYPE_CHECKING, Optional
+from typing import TYPE_CHECKING, Any, Optional
 
 import torch
 
@@ -27,6 +27,9 @@ class CacheInitParams:
     attn_tp_cache_group: Optional[torch.distributed.ProcessGroup] = None
     pp_cache_group: Optional[torch.distributed.ProcessGroup] = None
     eviction_policy: str = "lru"
+    # Keyword arguments for the eviction policy's constructor; see the strategy
+    # classes in evict_policy.py for what each policy accepts.
+    eviction_policy_config: Optional[dict[str, Any]] = None
     disable_finished_insert: bool = False
 
     enable_metrics: bool = False

python/sglang/srt/mem_cache/kv_cache_builder.py

@@ -306,6 +306,7 @@ def build_kv_cache(
         attn_tp_cache_group=attn_tp_cpu_group,
         pp_cache_group=pp_group.cpu_group,
         eviction_policy=get_memory().radix_eviction_policy,
+        eviction_policy_config=get_memory().radix_eviction_policy_config,
         enable_metrics=enable_metrics,
         enable_kv_cache_events=enable_kv_cache_events,
         enable_session_radix_cache=get_memory().enable_session_radix_cache,

python/sglang/srt/mem_cache/radix_cache.py

@@ -346,7 +346,9 @@ class RadixCache(BasePrefixCache):
         else:
             self.device = torch.device("cpu")
 
-        self.eviction_strategy = get_eviction_strategy(self.eviction_policy)
+        self.eviction_strategy = get_eviction_strategy(
+            self.eviction_policy, params.eviction_policy_config
+        )
 
         self.evictable_leaves = set()
         self.reset()

python/sglang/srt/mem_cache/rust_tree_core/adapter.py

@@ -327,6 +327,12 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
             raise ValueError(
                 "Rust TreeCore does not support component_registry_override"
             )
+        # The Rust core builds its own eviction strategy from the policy name
+        # alone, so a config would be dropped rather than applied.
+        if params.eviction_policy_config:
+            raise ValueError(
+                "Rust TreeCore does not support --radix-eviction-policy-config"
+            )
 
         self._page_size = params.page_size
         self.is_eagle = (

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

@@ -404,7 +404,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
         self.is_write_back = False
         self.has_swa_host_pool = False
         self.enable_session_radix_cache = params.enable_session_radix_cache
-        self.eviction_strategy = get_eviction_strategy(params.eviction_policy.lower())
+        self.eviction_strategy = get_eviction_strategy(
+            params.eviction_policy.lower(), params.eviction_policy_config
+        )
 
         # ``device`` is derived from the construction-time allocator; the
         # allocator/pool themselves are owned by the cache, not the tree.

python/sglang/srt/mem_cache/utils.py

@@ -56,7 +56,7 @@ from sglang.srt.mem_cache.evict_policy import (
     SLRUStrategy,
 )
 
-_EVICTION_POLICY_FACTORIES: dict[str, Callable[[], EvictionStrategy]] = {
+_EVICTION_POLICY_FACTORIES: dict[str, Callable[..., EvictionStrategy]] = {
     "lru": LRUStrategy,
     "lfu": LFUStrategy,
     "fifo": FIFOStrategy,
@@ -67,15 +67,19 @@ _EVICTION_POLICY_FACTORIES: dict[str, Callable[[], EvictionStrategy]] = {
 }
 
 
-def get_eviction_strategy(eviction_policy: str) -> EvictionStrategy:
+def get_eviction_strategy(
+    eviction_policy: str, config: Optional[dict[str, Any]] = None
+) -> EvictionStrategy:
+    """Build the eviction strategy; ``config`` is passed to it as keyword arguments."""
     policy = eviction_policy.lower()
     try:
-        return _EVICTION_POLICY_FACTORIES[policy]()
+        factory = _EVICTION_POLICY_FACTORIES[policy]
     except KeyError:
         supported = "', '".join(_EVICTION_POLICY_FACTORIES)
         raise ValueError(
             f"Unknown eviction policy: {policy}. Supported policies: '{supported}'."
         ) from None
+    return factory(**config) if config else factory()
 
 
 def maybe_init_custom_mem_pool(

python/sglang/srt/server_args.py

@@ -839,12 +839,30 @@ class ServerArgs:
                 "The eviction policy of radix trees. 'lru' stands for Least "
                 "Recently Used, 'lfu' stands for Least Frequently Used, 'slru' "
                 "stands for Segmented Least Recently Used, and 'priority' evicts "
-                "lower-priority requests first."
+                "lower-priority requests first. See "
+                "https://docs.sglang.io/docs/advanced_features/radix_eviction_policy "
+                "for what each policy optimizes for."
             ),
             choices=RADIX_EVICTION_POLICY_CHOICES,
         ),
         NS("memory"),
     ] = "lru"
+    radix_eviction_policy_config: A[
+        Optional[Dict[str, Any]],
+        Arg(
+            help=(
+                "Tuning parameters for --radix-eviction-policy, as a json object "
+                "passed to the policy as keyword arguments. Only 'slru' takes any "
+                "today: protected_threshold (int, default 2), e.g. "
+                "'{\"protected_threshold\": 4}'. An unrecognized key fails at "
+                "startup, naming the key and the policy. See "
+                "https://docs.sglang.io/docs/advanced_features/radix_eviction_policy#policy-parameters "
+                "for the full parameter list."
+            ),
+            type_parser=json.loads,
+        ),
+        NS("memory"),
+    ] = None
     prefill_only_disable_kv_cache: A[
         bool,
         "Skip the physical KV cache allocation for embedding-mode prefill-only workloads. Currently only valid with --is-embedding, --chunked-prefill-size=-1, --disable-radix-cache, an FA prefill backend, and non-FP4 KV cache so the fa_skip_kv_cache path is active (no layer reads or writes the cache). Other prefill-only workloads such as scoring/MIS may benefit from this later once their attention paths stop using paged KV. Scheduler admission accounting is unchanged; per-layer K/V tensors are sized to (page_size, head_num, head_dim) placeholders so GPU memory is not wasted.",