diff --git a/docs/advanced_features/dp_dpa_smg_guide.md b/docs/advanced_features/dp_dpa_smg_guide.md index 9ec5df64856e..74aea79031e3 100644 --- a/docs/advanced_features/dp_dpa_smg_guide.md +++ b/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 \ diff --git a/docs/advanced_features/sgl_model_gateway.md b/docs/advanced_features/sgl_model_gateway.md index 0f2da5b4776d..92b421466d6e 100644 --- a/docs/advanced_features/sgl_model_gateway.md +++ b/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 | diff --git a/sgl-model-gateway/README.md b/sgl-model-gateway/README.md index bffeae39f16e..f2650e245ae4 100644 --- a/sgl-model-gateway/README.md +++ b/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. diff --git a/sgl-model-gateway/bindings/python/src/lib.rs b/sgl-model-gateway/bindings/python/src/lib.rs index f0378fbcd039..6a032047deba 100644 --- a/sgl-model-gateway/bindings/python/src/lib.rs +++ b/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, diff --git a/sgl-model-gateway/bindings/python/src/sglang_router/router.py b/sgl-model-gateway/bindings/python/src/sglang_router/router.py index fa051b5765e8..af89acfed51e 100644 --- a/sgl-model-gateway/bindings/python/src/sglang_router/router.py +++ b/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 diff --git a/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py b/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py index d02a75992fe0..e874a48337f7 100644 --- a/sgl-model-gateway/bindings/python/src/sglang_router/router_args.py +++ b/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, diff --git a/sgl-model-gateway/bindings/python/tests/test_arg_parser.py b/sgl-model-gateway/bindings/python/tests/test_arg_parser.py index 43c8a401aa73..417e8c5477f4 100644 --- a/sgl-model-gateway/bindings/python/tests/test_arg_parser.py +++ b/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, diff --git a/sgl-model-gateway/bindings/python/tests/test_pyo3_binding.py b/sgl-model-gateway/bindings/python/tests/test_pyo3_binding.py index feea8329f82a..b06dec1d817c 100644 --- a/sgl-model-gateway/bindings/python/tests/test_pyo3_binding.py +++ b/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, diff --git a/sgl-model-gateway/bindings/python/tests/test_router_config.py b/sgl-model-gateway/bindings/python/tests/test_router_config.py index 5ba91cb1d724..46ff8849869c 100644 --- a/sgl-model-gateway/bindings/python/tests/test_router_config.py +++ b/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.""" diff --git a/sgl-model-gateway/bindings/python/tests/test_validation.py b/sgl-model-gateway/bindings/python/tests/test_validation.py index 587cd9504437..9ecb73678495 100644 --- a/sgl-model-gateway/bindings/python/tests/test_validation.py +++ b/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 diff --git a/sgl-model-gateway/src/config/builder.rs b/sgl-model-gateway/src/config/builder.rs index 5d4891ca6097..db39dcd646d0 100644 --- a/sgl-model-gateway/src/config/builder.rs +++ b/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"), + } + } } diff --git a/sgl-model-gateway/src/config/types.rs b/sgl-model-gateway/src/config/types.rs index b19802642ce1..9bbd44ed68f9 100644 --- a/sgl-model-gateway/src/config/types.rs +++ b/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, diff --git a/sgl-model-gateway/src/config/validation.rs b/sgl-model-gateway/src/config/validation.rs index 534aa8a4f4cd..be804b3c2d4b 100644 --- a/sgl-model-gateway/src/config/validation.rs +++ b/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, diff --git a/sgl-model-gateway/src/main.rs b/sgl-model-gateway/src/main.rs index b389bdad1944..3f092a107a72 100644 --- a/sgl-model-gateway/src/main.rs +++ b/sgl-model-gateway/src/main.rs @@ -154,6 +154,10 @@ struct CliArgs { #[arg(long, default_value_t = 0.3, help_heading = "Routing Policy")] cache_threshold: f32, + /// Weight for trading cache reuse against shadow cache balance in cache-aware routing + #[arg(long, default_value_t = 0.0, help_heading = "Routing Policy")] + cache_balance_weight: f32, + /// Absolute threshold for load balancing trigger #[arg(long, default_value_t = 64, help_heading = "Routing Policy")] balance_abs_threshold: usize, @@ -762,6 +766,7 @@ impl CliArgs { "round_robin" => PolicyConfig::RoundRobin, "cache_aware" => PolicyConfig::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, diff --git a/sgl-model-gateway/src/policies/cache_aware.rs b/sgl-model-gateway/src/policies/cache_aware.rs index 4a67d9dccd25..37a8ab353183 100644 --- a/sgl-model-gateway/src/policies/cache_aware.rs +++ b/sgl-model-gateway/src/policies/cache_aware.rs @@ -25,9 +25,11 @@ Process: a. For each request, find the worker with the highest prefix match b. If match rate > cache_threshold: - Route to the worker with highest match (likely has relevant data cached) + Route to the worker with highest match (likely has relevant data cached). When + cache_balance_weight > 0, subtract a shadow-cache-size penalty so common prefixes + can spread across workers instead of collapsing onto the first worker that saw them. c. If match rate ≤ cache_threshold: - Route to the worker with smallest tree size (most available cache capacity) + Route to the worker with the shortest queue. d. Background maintenance: Periodically evict least recently used leaf nodes to prevent memory overflow @@ -57,9 +59,14 @@ 5. max_tree_size: (integer) Maximum nodes per tree. When exceeded, LRU leaf nodes are evicted during the next eviction cycle. + + 6. cache_balance_weight: (float, >= 0.0) + Penalty applied to workers with larger shadow cache footprints when cache + hits are available. 0.0 preserves strict best-prefix routing. Higher values + trade some cache reuse for more even DP/cache distribution. */ -use std::sync::Arc; +use std::{cmp::Ordering, sync::Arc}; use async_trait::async_trait; use dashmap::DashMap; @@ -325,11 +332,8 @@ impl CacheAwarePolicy { ); } - // Use shortest queue when imbalanced - let min_load_idx = healthy_indices - .iter() - .min_by_key(|&&idx| workers[idx].load()) - .copied()?; + // Use shortest queue when imbalanced. + let min_load_idx = Self::select_min_load(workers, healthy_indices)?; // Even in imbalanced mode, update the tree to maintain cache state if let Some(text) = request_text { @@ -369,6 +373,62 @@ impl CacheAwarePolicy { Some(min_load_idx) } + + fn select_min_load( + workers: &[Arc], + healthy_indices: &[usize], + ) -> Option { + healthy_indices + .iter() + .min_by_key(|&&idx| workers[idx].load()) + .copied() + } + + fn select_cache_balanced_worker( + &self, + tree: &Tree, + text: &str, + workers: &[Arc], + healthy_indices: &[usize], + ) -> Option<(usize, usize)> { + let min_shadow_cache_chars = healthy_indices + .iter() + .map(|&idx| tree.tenant_char_count(workers[idx].url())) + .min()?; + + let mut best_idx = None; + let mut best_score = f64::NEG_INFINITY; + let mut best_match_chars = 0usize; + let mut best_shadow_cache_chars = usize::MAX; + let mut best_load = usize::MAX; + + for &idx in healthy_indices { + let worker = &workers[idx]; + let matched_chars = tree.prefix_match_tenant_count(text, worker.url()); + let shadow_cache_chars = tree.tenant_char_count(worker.url()); + let balance_penalty = self.config.cache_balance_weight as f64 + * shadow_cache_chars.saturating_sub(min_shadow_cache_chars) as f64; + let score = matched_chars as f64 - balance_penalty; + let load = worker.load(); + + let score_order = score.total_cmp(&best_score); + let is_better = best_idx.is_none() + || score_order == Ordering::Greater + || (score_order == Ordering::Equal + && (shadow_cache_chars, load, idx) + < (best_shadow_cache_chars, best_load, best_idx.unwrap())); + + if is_better { + best_idx = Some(idx); + best_score = score; + best_shadow_cache_chars = shadow_cache_chars; + best_load = load; + } + best_match_chars = best_match_chars.max(matched_chars); + } + + best_idx.map(|idx| (idx, best_match_chars)) + } } #[async_trait] @@ -420,7 +480,7 @@ impl LoadBalancingPolicy for CacheAwarePolicy { let tree = self.trees.get(&tree_key).map(|entry| entry.value().clone()); if let Some(tree) = tree { - // Now we work with the tree without holding the HashMap lock + // Now we work with the tree without holding the HashMap lock. // Use prefix_match_with_counts to avoid redundant chars().count() calls let result = tree.prefix_match_with_counts(text); let match_rate = if result.input_char_count == 0 { @@ -429,8 +489,21 @@ impl LoadBalancingPolicy for CacheAwarePolicy { result.matched_char_count as f32 / result.input_char_count as f32 }; - // Select worker without String allocation - let selected_idx = if match_rate > self.config.cache_threshold { + let selected_idx = if self.config.cache_balance_weight > 0.0 { + self.select_cache_balanced_worker(&tree, text, workers, &healthy_indices) + .and_then(|(idx, best_match_chars)| { + let match_rate = if result.input_char_count == 0 { + 0.0 + } else { + best_match_chars as f32 / result.input_char_count as f32 + }; + if match_rate > self.config.cache_threshold { + Some(idx) + } else { + Self::select_min_load(workers, &healthy_indices) + } + }) + } else if match_rate > self.config.cache_threshold { // Cache hit path: find worker by URL (compare &str directly, no allocation) let tenant_url: &str = &result.tenant; workers @@ -438,11 +511,8 @@ impl LoadBalancingPolicy for CacheAwarePolicy { .position(|w| w.url() == tenant_url) .filter(|&idx| workers[idx].is_healthy()) } else { - // Low cache match: use worker with minimum load - healthy_indices - .iter() - .min_by_key(|&&idx| workers[idx].load()) - .copied() + // Low cache match: use worker with minimum load. + Self::select_min_load(workers, &healthy_indices) }; if let Some(idx) = selected_idx { @@ -605,10 +675,107 @@ mod tests { assert_eq!(idx1, idx3); } + #[tokio::test] + async fn test_cache_aware_balance_weight_spreads_common_prefix() { + let policy = CacheAwarePolicy::with_config(CacheAwareConfig { + cache_threshold: 0.1, + cache_balance_weight: 1.0, + eviction_interval_secs: 0, + ..Default::default() + }); + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w3:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + ]; + policy.init_workers(&workers); + + let common_prefix = "shared system prompt ".repeat(20); + let mut selected = Vec::new(); + for suffix in ["conversation-a", "conversation-b", "conversation-c"] { + let prompt = format!("{common_prefix}{suffix}"); + selected.push( + policy + .select_worker( + &workers, + &SelectWorkerInfo { + request_text: Some(&prompt), + ..Default::default() + }, + ) + .await + .unwrap(), + ); + } + + assert_eq!(selected, vec![0, 1, 2]); + } + + #[tokio::test] + async fn test_cache_aware_balance_weight_preserves_strong_cache_hit() { + let policy = CacheAwarePolicy::with_config(CacheAwareConfig { + cache_threshold: 0.1, + cache_balance_weight: 0.2, + eviction_interval_secs: 0, + ..Default::default() + }); + let workers: Vec> = vec![ + Arc::new( + BasicWorkerBuilder::new("http://w1:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + Arc::new( + BasicWorkerBuilder::new("http://w2:8000") + .worker_type(WorkerType::Regular) + .build(), + ), + ]; + policy.init_workers(&workers); + + let prompt = "large reusable conversation prefix ".repeat(50); + let first = policy + .select_worker( + &workers, + &SelectWorkerInfo { + request_text: Some(&prompt), + ..Default::default() + }, + ) + .await + .unwrap(); + let second = policy + .select_worker( + &workers, + &SelectWorkerInfo { + request_text: Some(&prompt), + ..Default::default() + }, + ) + .await + .unwrap(); + + assert_eq!(first, 0); + assert_eq!(second, first); + } + #[tokio::test] async fn test_cache_aware_with_imbalanced_load() { let policy = CacheAwarePolicy::with_config(CacheAwareConfig { cache_threshold: 0.5, + cache_balance_weight: 0.0, balance_abs_threshold: 5, balance_rel_threshold: 2.0, eviction_interval_secs: 0, // Disable eviction thread @@ -958,10 +1125,10 @@ mod tests { decode_policy.init_workers(&decode_workers); let turns = [ - "turn1", "turn1 turn2", "turn1 turn2 turn3", "turn1 turn2 turn3 turn4", + "turn1 turn2 turn3 turn4 turn5", ]; let mut prefill_idx: Option = None; @@ -983,14 +1150,16 @@ mod tests { None => prefill_idx = Some(p), Some(pinned) => assert_eq!( p, pinned, - "prefill should stay pinned across turns (prompt={prompt:?})" + "prefill should stay pinned when reusable prefix exceeds threshold \ + (prompt={prompt:?})" ), } match decode_idx { None => decode_idx = Some(d), Some(pinned) => assert_eq!( d, pinned, - "decode should stay pinned across turns (prompt={prompt:?})" + "decode should stay pinned when reusable prefix exceeds threshold \ + (prompt={prompt:?})" ), } } @@ -1025,10 +1194,10 @@ mod tests { policy.init_workers(&combined); let turns = [ - "turn1", "turn1 turn2", "turn1 turn2 turn3", "turn1 turn2 turn3 turn4", + "turn1 turn2 turn3 turn4 turn5", ]; let mut prefill_idx: Option = None; @@ -1062,14 +1231,16 @@ mod tests { None => prefill_idx = Some(p), Some(pinned) => assert_eq!( p, pinned, - "prefill should stay pinned across turns (prompt={prompt:?})" + "prefill should stay pinned when reusable prefix exceeds threshold \ + (prompt={prompt:?})" ), } match decode_idx { None => decode_idx = Some(d), Some(pinned) => assert_eq!( d, pinned, - "decode should stay pinned across turns (prompt={prompt:?})" + "decode should stay pinned when reusable prefix exceeds threshold \ + (prompt={prompt:?})" ), } } diff --git a/sgl-model-gateway/src/policies/factory.rs b/sgl-model-gateway/src/policies/factory.rs index ae3ad845e5ae..24d947a0b49c 100644 --- a/sgl-model-gateway/src/policies/factory.rs +++ b/sgl-model-gateway/src/policies/factory.rs @@ -21,6 +21,7 @@ impl PolicyFactory { PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()), PolicyConfig::CacheAware { cache_threshold, + cache_balance_weight, balance_abs_threshold, balance_rel_threshold, eviction_interval_secs, @@ -28,6 +29,7 @@ impl PolicyFactory { } => { let config = CacheAwareConfig { cache_threshold: *cache_threshold, + cache_balance_weight: *cache_balance_weight, balance_abs_threshold: *balance_abs_threshold, balance_rel_threshold: *balance_rel_threshold, eviction_interval_secs: *eviction_interval_secs, @@ -110,6 +112,7 @@ mod tests { let policy = PolicyFactory::create_from_config(&PolicyConfig::CacheAware { cache_threshold: 0.7, + cache_balance_weight: 0.0, balance_abs_threshold: 10, balance_rel_threshold: 1.5, eviction_interval_secs: 30, diff --git a/sgl-model-gateway/src/policies/mod.rs b/sgl-model-gateway/src/policies/mod.rs index debc385a351b..5a8d186f6717 100644 --- a/sgl-model-gateway/src/policies/mod.rs +++ b/sgl-model-gateway/src/policies/mod.rs @@ -97,6 +97,7 @@ pub trait LoadBalancingPolicy: Send + Sync + Debug { #[derive(Debug, Clone)] pub struct CacheAwareConfig { pub cache_threshold: f32, + pub cache_balance_weight: f32, pub balance_abs_threshold: usize, pub balance_rel_threshold: f32, pub eviction_interval_secs: u64, @@ -107,6 +108,7 @@ impl Default for CacheAwareConfig { fn default() -> Self { Self { cache_threshold: 0.5, + cache_balance_weight: 0.0, balance_abs_threshold: 32, balance_rel_threshold: 1.1, eviction_interval_secs: 30, diff --git a/sgl-model-gateway/src/policies/tree.rs b/sgl-model-gateway/src/policies/tree.rs index 5a149ed87a2b..0e2d98c1e24f 100644 --- a/sgl-model-gateway/src/policies/tree.rs +++ b/sgl-model-gateway/src/policies/tree.rs @@ -628,6 +628,10 @@ impl Tree { #[allow(dead_code)] pub fn prefix_match_tenant(&self, text: &str, tenant: &str) -> String { + take_chars(text, self.prefix_match_tenant_count(text, tenant)) + } + + pub fn prefix_match_tenant_count(&self, text: &str, tenant: &str) -> usize { // Use slice-based traversal - no Vec allocation // Intern tenant ID once for efficient lookups @@ -688,8 +692,14 @@ impl Tree { .insert(Arc::clone(&tenant_id), epoch); } - // Build result from original input using char count - take_chars(text, matched_chars) + matched_chars + } + + pub fn tenant_char_count(&self, tenant: &str) -> usize { + self.tenant_char_count + .get(tenant) + .map(|count| *count) + .unwrap_or(0) } /// Return the list of tenants for which this node is a leaf. diff --git a/sgl-model-gateway/tests/routing/cache_aware_backward_compat_test.rs b/sgl-model-gateway/tests/routing/cache_aware_backward_compat_test.rs index 8e83651856a3..859aa9467c63 100644 --- a/sgl-model-gateway/tests/routing/cache_aware_backward_compat_test.rs +++ b/sgl-model-gateway/tests/routing/cache_aware_backward_compat_test.rs @@ -9,6 +9,7 @@ use smg::{ async fn test_backward_compatibility_with_empty_model_id() { let config = CacheAwareConfig { cache_threshold: 0.5, + cache_balance_weight: 0.0, balance_abs_threshold: 2, balance_rel_threshold: 1.5, eviction_interval_secs: 0, // Disable background eviction for testing @@ -60,6 +61,7 @@ async fn test_backward_compatibility_with_empty_model_id() { async fn test_mixed_model_ids() { let config = CacheAwareConfig { cache_threshold: 0.5, + cache_balance_weight: 0.0, balance_abs_threshold: 2, balance_rel_threshold: 1.5, eviction_interval_secs: 0, diff --git a/sgl-model-gateway/tests/routing/test_pd_routing.rs b/sgl-model-gateway/tests/routing/test_pd_routing.rs index 3587512b825d..aa1ef518fb6e 100644 --- a/sgl-model-gateway/tests/routing/test_pd_routing.rs +++ b/sgl-model-gateway/tests/routing/test_pd_routing.rs @@ -166,6 +166,7 @@ mod pd_routing_unit_tests { }, PolicyConfig::CacheAware { cache_threshold: 0.7, + cache_balance_weight: 0.0, balance_abs_threshold: 20, balance_rel_threshold: 1.2, eviction_interval_secs: 60,