Add cache balance weight for cache-aware routing

docs/advanced_features/dp_dpa_smg_guide.md

@@ -292,6 +292,7 @@ python -m sglang_router.launch_router \
     --worker-urls http://worker1:8000 http://worker2:8000 \
     --policy cache_aware \
     --cache-threshold 0.5 \
+    --cache-balance-weight 0.0 \
     --balance-abs-threshold 32 \
     --balance-rel-threshold 1.5 \
     --eviction-interval-secs 120 \

docs/advanced_features/sgl_model_gateway.md

@@ -1406,6 +1406,7 @@ python -m sglang_router.launch_router \
 | Parameter | Recommendation | Reason |
 |-----------|---------------|--------|
 | `--policy` | `cache_aware` | Best for repeated prompts, ~30% latency reduction |
+| `--cache-balance-weight` | `0.0` | Trade some cache reuse for more even shadow-cache distribution |
 | `--max-concurrent-requests` | 2-4x worker count | Prevent overload while maximizing throughput |
 | `--queue-size` | 2x max-concurrent | Buffer for burst traffic |
 | `--request-timeout-secs` | Based on max generation length | Prevent stuck requests |

sgl-model-gateway/README.md

@@ -736,7 +736,7 @@ Router flags map to these values:
 ## Load Balancing Policies
 - `random`: uniform random worker selection.
 - `round_robin`: sequential rotation with atomic counters.
-- `cache_aware`: maintains a prefix tree of prompts to route repeat traffic and evens load with configurable thresholds (`--cache-threshold`, `--balance-abs-threshold`, `--balance-rel-threshold`, `--eviction-interval`, `--max-tree-size`).
+- `cache_aware`: maintains a prefix tree of prompts to route repeat traffic and evens load with configurable thresholds (`--cache-threshold`, `--cache-balance-weight`, `--balance-abs-threshold`, `--balance-rel-threshold`, `--eviction-interval`, `--max-tree-size`).
 - `power_of_two`: chooses the lighter worker among two random candidates; integrates with `LoadMonitor`.
   Per-model overrides are available in PD mode (`--prefill-policy`, `--decode-policy`) and IGW mode via the worker registry.
 

sgl-model-gateway/bindings/python/src/lib.rs

@@ -352,6 +352,7 @@ struct Router {
     worker_startup_timeout_secs: u64,
     worker_startup_check_interval: u64,
     cache_threshold: f32,
+    cache_balance_weight: f32,
     balance_abs_threshold: usize,
     balance_rel_threshold: f32,
     eviction_interval_secs: u64,
@@ -467,6 +468,7 @@ impl Router {
                 PolicyType::RoundRobin => ConfigPolicyConfig::RoundRobin,
                 PolicyType::CacheAware => ConfigPolicyConfig::CacheAware {
                     cache_threshold: self.cache_threshold,
+                    cache_balance_weight: self.cache_balance_weight,
                     balance_abs_threshold: self.balance_abs_threshold,
                     balance_rel_threshold: self.balance_rel_threshold,
                     eviction_interval_secs: self.eviction_interval_secs,
@@ -675,6 +677,7 @@ impl Router {
         worker_startup_timeout_secs = 600,
         worker_startup_check_interval = 30,
         cache_threshold = 0.3,
+        cache_balance_weight = 0.0,
         balance_abs_threshold = 64,
         balance_rel_threshold = 1.5,
         eviction_interval_secs = 120,
@@ -767,6 +770,7 @@ impl Router {
         worker_startup_timeout_secs: u64,
         worker_startup_check_interval: u64,
         cache_threshold: f32,
+        cache_balance_weight: f32,
         balance_abs_threshold: usize,
         balance_rel_threshold: f32,
         eviction_interval_secs: u64,
@@ -872,6 +876,7 @@ impl Router {
             worker_startup_timeout_secs,
             worker_startup_check_interval,
             cache_threshold,
+            cache_balance_weight,
             balance_abs_threshold,
             balance_rel_threshold,
             eviction_interval_secs,

sgl-model-gateway/bindings/python/src/sglang_router/router.py

@@ -161,8 +161,10 @@ class Router:
         worker_startup_timeout_secs: Timeout in seconds for worker startup and registration. Large models can take significant time to load into GPU memory. Default: 1800 (30 minutes)
         worker_startup_check_interval: Interval in seconds between checks for worker initialization. Default: 10
         cache_threshold: Cache threshold (0.0-1.0) for cache-aware routing. Routes to cached worker
-            if the match rate exceeds threshold, otherwise routes to the worker with the smallest
-            tree. Default: 0.5
+            if the match rate exceeds threshold, otherwise routes to the worker with the shortest
+            queue. Default: 0.5
+        cache_balance_weight: Penalty weight for trading cache reuse against shadow-cache balance
+            in cache-aware routing. 0.0 preserves strict best-prefix routing. Default: 0.0
         balance_abs_threshold: Load balancing is triggered when (max_load - min_load) > abs_threshold
             AND max_load > min_load * rel_threshold. Otherwise, use cache aware. Default: 32
         balance_rel_threshold: Load balancing is triggered when (max_load - min_load) > abs_threshold

sgl-model-gateway/bindings/python/src/sglang_router/router_args.py

@@ -57,6 +57,7 @@ class RouterArgs:
     worker_startup_timeout_secs: int = 1800
     worker_startup_check_interval: int = 30
     cache_threshold: float = 0.3
+    cache_balance_weight: float = 0.0
     balance_abs_threshold: int = 64
     balance_rel_threshold: float = 1.5
     eviction_interval_secs: int = 60
@@ -307,6 +308,12 @@ def add_cli_args(
             default=RouterArgs.cache_threshold,
             help="Cache threshold (0.0-1.0) for cache-aware routing",
         )
+        routing_group.add_argument(
+            f"--{prefix}cache-balance-weight",
+            type=float,
+            default=RouterArgs.cache_balance_weight,
+            help="Penalty weight for trading cache reuse against shadow-cache balance in cache-aware routing",
+        )
         routing_group.add_argument(
             f"--{prefix}balance-abs-threshold",
             type=int,

sgl-model-gateway/bindings/python/tests/test_arg_parser.py

@@ -151,6 +151,7 @@ def test_from_cli_args_basic(self):
             router_worker_startup_timeout_secs=300,
             router_worker_startup_check_interval=15,
             router_cache_threshold=0.7,
+            router_cache_balance_weight=0.25,
             router_balance_abs_threshold=128,
             router_balance_rel_threshold=2.0,
             router_eviction_interval=180,
@@ -278,6 +279,7 @@ def test_from_cli_args_pd_mode(self):
             router_worker_startup_timeout_secs=600,
             router_worker_startup_check_interval=30,
             router_cache_threshold=0.3,
+            router_cache_balance_weight=0.0,
             router_balance_abs_threshold=64,
             router_balance_rel_threshold=1.5,
             router_eviction_interval=120,

sgl-model-gateway/bindings/python/tests/test_pyo3_binding.py

@@ -384,6 +384,7 @@ def test_kitchen_sink_passes_every_field(self):
             worker_startup_timeout_secs=60,
             worker_startup_check_interval=5,
             cache_threshold=0.5,
+            cache_balance_weight=0.25,
             balance_abs_threshold=32,
             balance_rel_threshold=1.2,
             eviction_interval_secs=30,

sgl-model-gateway/bindings/python/tests/test_router_config.py

@@ -369,6 +369,7 @@ def test_config_serialization(self):
             worker_urls=["http://worker1:8000"],
             policy="cache_aware",
             cache_threshold=0.5,
+            cache_balance_weight=0.25,
         )
 
         # Test that we can access all attributes
@@ -377,6 +378,7 @@ def test_config_serialization(self):
         assert hasattr(args, "worker_urls")
         assert hasattr(args, "policy")
         assert hasattr(args, "cache_threshold")
+        assert hasattr(args, "cache_balance_weight")
 
     def test_config_with_none_values(self):
         """Test configuration with None values."""

sgl-model-gateway/bindings/python/tests/test_validation.py

@@ -121,6 +121,12 @@ def test_cache_threshold_validation(self):
             args = RouterArgs(cache_threshold=threshold)
             assert args.cache_threshold == threshold
 
+    def test_cache_balance_weight_validation(self):
+        """Test cache balance weight parameter plumbing."""
+        for weight in [0.0, 0.1, 1.0]:
+            args = RouterArgs(cache_balance_weight=weight)
+            assert args.cache_balance_weight == weight
+
     def test_balance_threshold_validation(self):
         """Test load balancing threshold parameter validation."""
         # Valid absolute thresholds

sgl-model-gateway/src/config/builder.rs

@@ -118,6 +118,27 @@ impl RouterConfigBuilder {
     ) -> Self {
         self.config.policy = PolicyConfig::CacheAware {
             cache_threshold,
+            cache_balance_weight: 0.0,
+            balance_abs_threshold,
+            balance_rel_threshold,
+            eviction_interval_secs,
+            max_tree_size,
+        };
+        self
+    }
+
+    pub fn cache_aware_policy_with_balance(
+        mut self,
+        cache_threshold: f32,
+        cache_balance_weight: f32,
+        balance_abs_threshold: usize,
+        balance_rel_threshold: f32,
+        eviction_interval_secs: u64,
+        max_tree_size: usize,
+    ) -> Self {
+        self.config.policy = PolicyConfig::CacheAware {
+            cache_threshold,
+            cache_balance_weight,
             balance_abs_threshold,
             balance_rel_threshold,
             eviction_interval_secs,
@@ -863,4 +884,26 @@ mod tests {
             _ => panic!("Expected CacheAware policy"),
         }
     }
+
+    /// Test cache-aware helper with cache balance weight
+    #[test]
+    fn test_builder_cache_aware_policy_with_balance() {
+        let config = RouterConfigBuilder::new()
+            .regular_mode(vec!["http://worker1:8000".to_string()])
+            .cache_aware_policy_with_balance(0.8, 0.25, 10, 1.5, 300, 1000)
+            .build()
+            .unwrap();
+
+        match config.policy {
+            PolicyConfig::CacheAware {
+                cache_threshold,
+                cache_balance_weight,
+                ..
+            } => {
+                assert!((cache_threshold - 0.8).abs() < 0.0001);
+                assert!((cache_balance_weight - 0.25).abs() < 0.0001);
+            }
+            _ => panic!("Expected CacheAware policy"),
+        }
+    }
 }

sgl-model-gateway/src/config/types.rs

@@ -261,6 +261,8 @@ pub enum PolicyConfig {
     #[serde(rename = "cache_aware")]
     CacheAware {
         cache_threshold: f32,
+        #[serde(default)]
+        cache_balance_weight: f32,
         balance_abs_threshold: usize,
         balance_rel_threshold: f32,
         eviction_interval_secs: u64,
@@ -808,6 +810,7 @@ mod tests {
 
         let cache_aware = PolicyConfig::CacheAware {
             cache_threshold: 0.8,
+            cache_balance_weight: 0.0,
             balance_abs_threshold: 10,
             balance_rel_threshold: 1.5,
             eviction_interval_secs: 300,
@@ -829,6 +832,7 @@ mod tests {
 
         let cache_aware = PolicyConfig::CacheAware {
             cache_threshold: 0.8,
+            cache_balance_weight: 0.0,
             balance_abs_threshold: 10,
             balance_rel_threshold: 1.5,
             eviction_interval_secs: 300,
@@ -837,6 +841,7 @@ mod tests {
         let json = serde_json::to_string(&cache_aware).unwrap();
         assert!(json.contains("\"type\":\"cache_aware\""));
         assert!(json.contains("\"cache_threshold\":0.8"));
+        assert!(json.contains("\"cache_balance_weight\":0.0"));
         assert!(json.contains("\"balance_abs_threshold\":10"));
 
         let power_of_two = PolicyConfig::PowerOfTwo {
@@ -851,6 +856,7 @@ mod tests {
     fn test_cache_aware_parameters() {
         let cache_aware = PolicyConfig::CacheAware {
             cache_threshold: 0.75,
+            cache_balance_weight: 0.25,
             balance_abs_threshold: 20,
             balance_rel_threshold: 2.0,
             eviction_interval_secs: 600,
@@ -860,12 +866,14 @@ mod tests {
         match cache_aware {
             PolicyConfig::CacheAware {
                 cache_threshold,
+                cache_balance_weight,
                 balance_abs_threshold,
                 balance_rel_threshold,
                 eviction_interval_secs,
                 max_tree_size,
             } => {
                 assert!((cache_threshold - 0.75).abs() < 0.0001);
+                assert!((cache_balance_weight - 0.25).abs() < 0.0001);
                 assert_eq!(balance_abs_threshold, 20);
                 assert!((balance_rel_threshold - 2.0).abs() < 0.0001);
                 assert_eq!(eviction_interval_secs, 600);
@@ -875,6 +883,29 @@ mod tests {
         }
     }
 
+    #[test]
+    fn test_cache_aware_deserialize_default_cache_balance_weight() {
+        let json = r#"{
+            "type": "cache_aware",
+            "cache_threshold": 0.8,
+            "balance_abs_threshold": 10,
+            "balance_rel_threshold": 1.5,
+            "eviction_interval_secs": 300,
+            "max_tree_size": 1000
+        }"#;
+
+        let cache_aware: PolicyConfig = serde_json::from_str(json).unwrap();
+        match cache_aware {
+            PolicyConfig::CacheAware {
+                cache_balance_weight,
+                ..
+            } => {
+                assert_eq!(cache_balance_weight, 0.0);
+            }
+            _ => panic!("Expected CacheAware"),
+        }
+    }
+
     #[test]
     fn test_power_of_two_parameters() {
         let power_of_two = PolicyConfig::PowerOfTwo {
@@ -1253,6 +1284,7 @@ mod tests {
             decode_urls: vec!["http://decode1".to_string()],
             prefill_policy: Some(PolicyConfig::CacheAware {
                 cache_threshold: 0.5,
+                cache_balance_weight: 0.0,
                 balance_abs_threshold: 32,
                 balance_rel_threshold: 1.1,
                 eviction_interval_secs: 60,
@@ -1283,6 +1315,7 @@ mod tests {
             decode_urls: vec!["http://decode1".to_string()],
             prefill_policy: Some(PolicyConfig::CacheAware {
                 cache_threshold: 0.5,
+                cache_balance_weight: 0.0,
                 balance_abs_threshold: 32,
                 balance_rel_threshold: 1.1,
                 eviction_interval_secs: 60,
@@ -1339,6 +1372,7 @@ mod tests {
 
         let main_policy = PolicyConfig::CacheAware {
             cache_threshold: 0.7,
+            cache_balance_weight: 0.0,
             balance_abs_threshold: 20,
             balance_rel_threshold: 1.5,
             eviction_interval_secs: 300,

sgl-model-gateway/src/config/validation.rs

@@ -153,6 +153,7 @@ impl ConfigValidator {
             | PolicyConfig::ConsistentHashing => {}
             PolicyConfig::CacheAware {
                 cache_threshold,
+                cache_balance_weight,
                 balance_abs_threshold: _,
                 balance_rel_threshold,
                 eviction_interval_secs,
@@ -166,6 +167,14 @@ impl ConfigValidator {
                     });
                 }
 
+                if !cache_balance_weight.is_finite() || *cache_balance_weight < 0.0 {
+                    return Err(ConfigError::InvalidValue {
+                        field: "cache_balance_weight".to_string(),
+                        value: cache_balance_weight.to_string(),
+                        reason: "Must be finite and >= 0.0".to_string(),
+                    });
+                }
+
                 if *balance_rel_threshold < 1.0 {
                     return Err(ConfigError::InvalidValue {
                         field: "balance_rel_threshold".to_string(),
@@ -725,6 +734,29 @@ mod tests {
             },
             PolicyConfig::CacheAware {
                 cache_threshold: 1.5, // Invalid: > 1.0
+                cache_balance_weight: 0.0,
+                balance_abs_threshold: 32,
+                balance_rel_threshold: 1.1,
+                eviction_interval_secs: 60,
+                max_tree_size: 1000,
+            },
+        );
+
+        assert!(ConfigValidator::validate(&config).is_err());
+    }
+
+    #[test]
+    fn test_validate_cache_aware_balance_weight() {
+        let config = RouterConfig::new(
+            RoutingMode::Regular {
+                worker_urls: vec![
+                    "http://worker1:8000".to_string(),
+                    "http://worker2:8000".to_string(),
+                ],
+            },
+            PolicyConfig::CacheAware {
+                cache_threshold: 0.5,
+                cache_balance_weight: -0.1,
                 balance_abs_threshold: 32,
                 balance_rel_threshold: 1.1,
                 eviction_interval_secs: 60,
@@ -744,6 +776,7 @@ mod tests {
             },
             PolicyConfig::CacheAware {
                 cache_threshold: 0.5,
+                cache_balance_weight: 0.0,
                 balance_abs_threshold: 32,
                 balance_rel_threshold: 1.1,
                 eviction_interval_secs: 60,
@@ -798,6 +831,7 @@ mod tests {
             },
             PolicyConfig::CacheAware {
                 cache_threshold: 0.5,
+                cache_balance_weight: 0.0,
                 balance_abs_threshold: 32,
                 balance_rel_threshold: 1.1,
                 eviction_interval_secs: 60,
@@ -842,6 +876,7 @@ mod tests {
                 ],
                 prefill_policy: Some(PolicyConfig::CacheAware {
                     cache_threshold: 0.5,
+                    cache_balance_weight: 0.0,
                     balance_abs_threshold: 32,
                     balance_rel_threshold: 1.1,
                     eviction_interval_secs: 60,