From 3645ed0f737c1d44fe187ea655d60784a964d8c6 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Sat, 27 Dec 2025 05:58:05 -0500 Subject: [PATCH] [model-gateway] Add PrefixHash load balancing policy for KV cache-aware routing (#15935) --- sgl-model-gateway/Cargo.toml | 2 + .../bindings/python/sglang_router/router.py | 1 + sgl-model-gateway/bindings/python/src/lib.rs | 5 + sgl-model-gateway/src/config/types.rs | 25 ++ sgl-model-gateway/src/config/validation.rs | 20 + sgl-model-gateway/src/main.rs | 18 +- .../src/observability/metrics.rs | 9 + sgl-model-gateway/src/policies/factory.rs | 14 +- sgl-model-gateway/src/policies/mod.rs | 5 + sgl-model-gateway/src/policies/prefix_hash.rs | 409 ++++++++++++++++++ sgl-model-gateway/src/policies/registry.rs | 62 +-- .../grpc/common/stages/worker_selection.rs | 20 +- .../src/routers/http/pd_router.rs | 1 + sgl-model-gateway/src/routers/http/router.rs | 1 + 14 files changed, 532 insertions(+), 60 deletions(-) create mode 100644 sgl-model-gateway/src/policies/prefix_hash.rs diff --git a/sgl-model-gateway/Cargo.toml b/sgl-model-gateway/Cargo.toml index eaa68340a..5cba00d7d 100644 --- a/sgl-model-gateway/Cargo.toml +++ b/sgl-model-gateway/Cargo.toml @@ -50,6 +50,8 @@ futures = "0.3" dashmap = "6.1.0" lru = "0.16.2" blake3 = "1.5" +xxhash-rust = { version = "0.8", features = ["xxh3"] } +bytemuck = { version = "1.21", features = ["derive"] } http = "1.1.0" tokio = { version = "1.42.0", features = ["full"] } async-trait = "0.1" diff --git a/sgl-model-gateway/bindings/python/sglang_router/router.py b/sgl-model-gateway/bindings/python/sglang_router/router.py index 95484cb2b..1f73915e0 100644 --- a/sgl-model-gateway/bindings/python/sglang_router/router.py +++ b/sgl-model-gateway/bindings/python/sglang_router/router.py @@ -27,6 +27,7 @@ def policy_from_str(policy_str: Optional[str]) -> PolicyType: "bucket": PolicyType.Bucket, "manual": PolicyType.Manual, "consistent_hashing": PolicyType.ConsistentHashing, + "prefix_hash": PolicyType.PrefixHash, } return policy_map[policy_str] diff --git a/sgl-model-gateway/bindings/python/src/lib.rs b/sgl-model-gateway/bindings/python/src/lib.rs index 2ea27573a..e66138ab3 100644 --- a/sgl-model-gateway/bindings/python/src/lib.rs +++ b/sgl-model-gateway/bindings/python/src/lib.rs @@ -14,6 +14,7 @@ pub enum PolicyType { Bucket, Manual, ConsistentHashing, + PrefixHash, } #[pyclass(eq)] @@ -418,6 +419,10 @@ impl Router { }, PolicyType::Manual => ConfigPolicyConfig::Manual, PolicyType::ConsistentHashing => ConfigPolicyConfig::ConsistentHashing, + PolicyType::PrefixHash => ConfigPolicyConfig::PrefixHash { + prefix_token_count: 256, + load_factor: 1.25, + }, } }; diff --git a/sgl-model-gateway/src/config/types.rs b/sgl-model-gateway/src/config/types.rs index fbc7bbc53..cbca5f614 100644 --- a/sgl-model-gateway/src/config/types.rs +++ b/sgl-model-gateway/src/config/types.rs @@ -350,6 +350,30 @@ pub enum PolicyConfig { /// - Provides O(log n) lookup with minimal redistribution (~1/N keys) on topology change #[serde(rename = "consistent_hashing")] ConsistentHashing, + + /// Prefix hash policy for KV cache-aware load balancing. + /// A lightweight alternative to cache_aware radix tree. + /// Routes requests based on prefix token hash for cache locality. + /// - Uses consistent hash ring with bounded load balancing + /// - Walks ring if worker is overloaded (load > avg * load_factor) + /// - O(log n) lookup instead of O(prefix_len) radix tree traversal + #[serde(rename = "prefix_hash")] + PrefixHash { + /// Number of prefix tokens to hash (default: 256) + #[serde(default = "default_prefix_token_count")] + prefix_token_count: usize, + /// Load factor threshold - walk ring if load > avg * factor (default: 1.25) + #[serde(default = "default_load_factor")] + load_factor: f64, + }, +} + +fn default_prefix_token_count() -> usize { + 256 +} + +fn default_load_factor() -> f64 { + 1.25 } impl PolicyConfig { @@ -362,6 +386,7 @@ impl PolicyConfig { PolicyConfig::Bucket { .. } => "bucket", PolicyConfig::Manual => "manual", PolicyConfig::ConsistentHashing => "consistent_hashing", + PolicyConfig::PrefixHash { .. } => "prefix_hash", } } } diff --git a/sgl-model-gateway/src/config/validation.rs b/sgl-model-gateway/src/config/validation.rs index 2209b1c7b..53c04c4d1 100644 --- a/sgl-model-gateway/src/config/validation.rs +++ b/sgl-model-gateway/src/config/validation.rs @@ -229,6 +229,26 @@ impl ConfigValidator { }); } } + PolicyConfig::PrefixHash { + prefix_token_count, + load_factor, + } => { + if *prefix_token_count == 0 { + return Err(ConfigError::InvalidValue { + field: "prefix_token_count".to_string(), + value: prefix_token_count.to_string(), + reason: "Must be > 0".to_string(), + }); + } + + if *load_factor < 1.0 { + return Err(ConfigError::InvalidValue { + field: "load_factor".to_string(), + value: load_factor.to_string(), + reason: "Must be >= 1.0".to_string(), + }); + } + } } Ok(()) } diff --git a/sgl-model-gateway/src/main.rs b/sgl-model-gateway/src/main.rs index 445a7f8cf..abbb102a6 100644 --- a/sgl-model-gateway/src/main.rs +++ b/sgl-model-gateway/src/main.rs @@ -137,7 +137,7 @@ struct CliArgs { #[arg(long, num_args = 0..)] worker_urls: Vec, - #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "manual"])] + #[arg(long, default_value = "cache_aware", value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "prefix_hash", "manual"])] policy: String, #[arg(long, default_value_t = false)] @@ -146,10 +146,10 @@ struct CliArgs { #[arg(long, action = ArgAction::Append)] decode: Vec, - #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "manual"])] + #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "prefix_hash", "manual"])] prefill_policy: Option, - #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "manual"])] + #[arg(long, value_parser = ["random", "round_robin", "cache_aware", "power_of_two", "prefix_hash", "manual"])] decode_policy: Option, #[arg(long, default_value_t = 1800)] @@ -173,6 +173,14 @@ struct CliArgs { #[arg(long, default_value_t = 67108864)] max_tree_size: usize, + /// Number of prefix tokens to use for prefix_hash policy (default: 256) + #[arg(long, default_value_t = 256)] + prefix_token_count: usize, + + /// Load factor threshold for prefix_hash policy (default: 1.25) + #[arg(long, default_value_t = 1.25)] + prefix_hash_load_factor: f64, + #[arg(long, default_value_t = 536870912)] max_payload_size: usize, @@ -549,6 +557,10 @@ impl CliArgs { "power_of_two" => PolicyConfig::PowerOfTwo { load_check_interval_secs: 5, }, + "prefix_hash" => PolicyConfig::PrefixHash { + prefix_token_count: self.prefix_token_count, + load_factor: self.prefix_hash_load_factor, + }, "manual" => PolicyConfig::Manual, _ => PolicyConfig::RoundRobin, } diff --git a/sgl-model-gateway/src/observability/metrics.rs b/sgl-model-gateway/src/observability/metrics.rs index 8f65390a6..5523f5c75 100644 --- a/sgl-model-gateway/src/observability/metrics.rs +++ b/sgl-model-gateway/src/observability/metrics.rs @@ -819,6 +819,15 @@ impl Metrics { .increment(1); } + /// Record prefix hash policy execution branch for routing decisions + pub fn record_worker_prefix_hash_policy_branch(branch: &'static str) { + counter!( + "smg_prefix_hash_policy_branch_total", + "branch" => branch + ) + .increment(1); + } + /// Set running requests per worker pub fn set_worker_requests_active(worker: &str, count: usize) { gauge!( diff --git a/sgl-model-gateway/src/policies/factory.rs b/sgl-model-gateway/src/policies/factory.rs index 3d7c9b4f0..2562303d0 100644 --- a/sgl-model-gateway/src/policies/factory.rs +++ b/sgl-model-gateway/src/policies/factory.rs @@ -4,7 +4,8 @@ use std::sync::Arc; use super::{ BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy, - LoadBalancingPolicy, ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, + LoadBalancingPolicy, ManualPolicy, PowerOfTwoPolicy, PrefixHashConfig, PrefixHashPolicy, + RandomPolicy, RoundRobinPolicy, }; use crate::config::PolicyConfig; @@ -48,6 +49,16 @@ impl PolicyFactory { } PolicyConfig::Manual => Arc::new(ManualPolicy::new()), PolicyConfig::ConsistentHashing => Arc::new(ConsistentHashingPolicy::new()), + PolicyConfig::PrefixHash { + prefix_token_count, + load_factor, + } => { + let config = PrefixHashConfig { + prefix_token_count: *prefix_token_count, + load_factor: *load_factor, + }; + Arc::new(PrefixHashPolicy::new(config)) + } } } @@ -63,6 +74,7 @@ impl PolicyFactory { "consistent_hashing" | "consistenthashing" => { Some(Arc::new(ConsistentHashingPolicy::new())) } + "prefix_hash" | "prefixhash" => Some(Arc::new(PrefixHashPolicy::with_defaults())), _ => None, } } diff --git a/sgl-model-gateway/src/policies/mod.rs b/sgl-model-gateway/src/policies/mod.rs index 48a908aca..709decf71 100644 --- a/sgl-model-gateway/src/policies/mod.rs +++ b/sgl-model-gateway/src/policies/mod.rs @@ -13,6 +13,7 @@ mod consistent_hashing; mod factory; mod manual; mod power_of_two; +mod prefix_hash; mod random; mod registry; mod round_robin; @@ -23,6 +24,7 @@ pub use consistent_hashing::ConsistentHashingPolicy; pub use factory::PolicyFactory; pub use manual::ManualPolicy; pub use power_of_two::PowerOfTwoPolicy; +pub use prefix_hash::{PrefixHashConfig, PrefixHashPolicy}; pub use random::RandomPolicy; pub use registry::PolicyRegistry; pub use round_robin::RoundRobinPolicy; @@ -144,6 +146,9 @@ pub(crate) fn normalize_model_key(model_id: &str) -> &str { pub struct SelectWorkerInfo<'a> { /// Request text for cache-aware routing pub request_text: Option<&'a str>, + /// Tokenized request for prefix-hash routing + /// Used by PrefixHashPolicy for token-based prefix hashing + pub tokens: Option<&'a [u32]>, /// HTTP headers for header-based routing policies /// Policies can extract routing information from headers like: /// - X-SMG-Target-Worker: Direct routing to a specific worker by index diff --git a/sgl-model-gateway/src/policies/prefix_hash.rs b/sgl-model-gateway/src/policies/prefix_hash.rs new file mode 100644 index 000000000..eeda785b0 --- /dev/null +++ b/sgl-model-gateway/src/policies/prefix_hash.rs @@ -0,0 +1,409 @@ +//! Prefix Hash routing policy for KV cache-aware load balancing +//! +//! A lightweight alternative to the full radix tree cache_aware policy. +//! Routes requests based on a hash of their prefix tokens to maximize +//! KV cache hits across workers. +//! +//! ## Algorithm +//! +//! 1. Extract first N tokens from the request (configurable prefix length) +//! 2. Hash the token sequence using xxhash for fast, stable hashing +//! 3. Use consistent hash ring to find the target worker +//! 4. If worker is overloaded (load > avg * load_factor), find least loaded +//! 5. Return least loaded worker that passes load check, or initial if all overloaded +//! +//! ## Complexity +//! +//! - Hash computation: O(prefix_length) +//! - Ring lookup: O(log n) binary search +//! - Load balance fallback: O(n) scan for least loaded +//! +//! ## Comparison with cache_aware +//! +//! | Aspect | prefix_hash | cache_aware (radix) | +//! |-----------------|-------------------|---------------------| +//! | Lookup | O(log n) | O(prefix_len) | +//! | Memory | O(workers × vn) | O(total_tokens) | +//! | Update | O(1) | O(prefix_len) | +//! | Precision | Prefix grouping | Exact matching | +//! +//! prefix_hash trades optimal cache utilization for predictable O(log n) performance. + +use std::sync::Arc; + +use super::{LoadBalancingPolicy, SelectWorkerInfo}; +use crate::{core::Worker, observability::metrics::Metrics}; + +/// Configuration for the PrefixHash load balancing policy +#[derive(Debug, Clone)] +pub struct PrefixHashConfig { + /// Number of prefix tokens to use for hashing. + /// Longer prefixes = more precise routing but less grouping. + /// Shorter prefixes = more requests grouped together. + /// Default: 256 tokens (~1 paragraph of text) + pub prefix_token_count: usize, + + /// Load factor threshold for walking the ring. + /// If a worker's load > (total_load / num_workers) * load_factor, + /// walk clockwise to the next worker. + /// Default: 1.25 (125% of average load) + pub load_factor: f64, +} + +impl Default for PrefixHashConfig { + fn default() -> Self { + Self { + prefix_token_count: 256, + load_factor: 1.25, + } + } +} + +/// Execution branch for metrics +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +enum Branch { + NoHealthyWorkers, + NoTokens, + RingHit, + LoadBalanceWalk, + FallbackLeastLoad, +} + +impl Branch { + #[inline] + const fn as_str(&self) -> &'static str { + match self { + Self::NoHealthyWorkers => "no_healthy_workers", + Self::NoTokens => "no_tokens", + Self::RingHit => "ring_hit", + Self::LoadBalanceWalk => "load_balance_walk", + Self::FallbackLeastLoad => "fallback_least_load", + } + } +} + +/// Prefix Hash load balancing policy +/// +/// Routes requests based on prefix token hash for KV cache locality. +/// Uses consistent hashing with bounded load balancing. +#[derive(Debug)] +pub struct PrefixHashPolicy { + config: PrefixHashConfig, +} + +impl PrefixHashPolicy { + /// Create a new PrefixHashPolicy with the given configuration + pub fn new(config: PrefixHashConfig) -> Self { + Self { config } + } + + /// Create a new PrefixHashPolicy with default configuration + pub fn with_defaults() -> Self { + Self::new(PrefixHashConfig::default()) + } + + /// Compute hash of prefix tokens using xxhash + #[inline] + fn compute_prefix_hash(&self, tokens: &[u32]) -> u64 { + let prefix_len = tokens.len().min(self.config.prefix_token_count); + let prefix = &tokens[..prefix_len]; + + let bytes: &[u8] = bytemuck::cast_slice(prefix); + xxhash_rust::xxh3::xxh3_64(bytes) + } + + /// Check if a worker's load is acceptable + #[inline] + fn load_ok(&self, worker_load: usize, total_load: usize, num_workers: usize) -> bool { + if total_load == 0 || num_workers == 0 { + return true; + } + + // Average load per worker (with +1 to simulate incoming request) + let avg_load = (total_load + 1) as f64 / num_workers as f64; + let threshold = avg_load * self.config.load_factor; + + (worker_load as f64) <= threshold + } + + /// Find worker using consistent hash ring with load balancing + fn find_worker_with_load_balance( + &self, + workers: &[Arc], + info: &SelectWorkerInfo, + prefix_hash: u64, + ) -> (Option, Branch) { + // Build healthy worker URL to index map + let healthy_workers: Vec<(usize, &Arc)> = workers + .iter() + .enumerate() + .filter(|(_, w)| w.is_healthy()) + .collect(); + + if healthy_workers.is_empty() { + return (None, Branch::NoHealthyWorkers); + } + + // Calculate total load for load balancing + let total_load: usize = healthy_workers.iter().map(|(_, w)| w.load()).sum(); + let num_workers = healthy_workers.len(); + + // Use pre-computed ring if available + if let Some(ref ring) = info.hash_ring { + // Convert prefix hash to a ring key string for lookup + let key = format!("{:016x}", prefix_hash); + + // Build URL to (index, worker) map for healthy workers + let healthy_url_map: std::collections::HashMap<&str, (usize, &Arc)> = + healthy_workers + .iter() + .map(|(idx, w)| (w.url(), (*idx, *w))) + .collect(); + + // Find initial worker from ring + if let Some(initial_url) = + ring.find_healthy_url(&key, |url| healthy_url_map.contains_key(url)) + { + if let Some(&(idx, worker)) = healthy_url_map.get(initial_url) { + let worker_load = worker.load(); + + // Check if initial worker has acceptable load + if self.load_ok(worker_load, total_load, num_workers) { + return (Some(idx), Branch::RingHit); + } + + // Initial worker overloaded, find least loaded healthy worker + // This is a simpler approach than walking the ring + let least_loaded = healthy_workers + .iter() + .filter(|(_, w)| self.load_ok(w.load(), total_load, num_workers)) + .min_by_key(|(_, w)| w.load()); + + if let Some(&(idx, _)) = least_loaded { + return (Some(idx), Branch::LoadBalanceWalk); + } + + // All workers overloaded, use initial worker anyway + return (Some(idx), Branch::LoadBalanceWalk); + } + } + } + + // Fallback: no ring or ring lookup failed, use least loaded worker + let least_loaded = healthy_workers + .iter() + .min_by_key(|(_, w)| w.load()) + .map(|(idx, _)| *idx); + + (least_loaded, Branch::FallbackLeastLoad) + } + + fn select_worker_impl( + &self, + workers: &[Arc], + info: &SelectWorkerInfo, + ) -> (Option, Branch) { + if workers.is_empty() { + return (None, Branch::NoHealthyWorkers); + } + + // Get tokens from SelectWorkerInfo + let tokens = match info.tokens { + Some(t) if !t.is_empty() => t, + _ => return (None, Branch::NoTokens), + }; + + // Compute prefix hash + let prefix_hash = self.compute_prefix_hash(tokens); + + // Find worker using ring with load balancing + self.find_worker_with_load_balance(workers, info, prefix_hash) + } +} + +impl LoadBalancingPolicy for PrefixHashPolicy { + fn select_worker(&self, workers: &[Arc], info: &SelectWorkerInfo) -> Option { + let (result, branch) = self.select_worker_impl(workers, info); + Metrics::record_worker_prefix_hash_policy_branch(branch.as_str()); + result + } + + fn name(&self) -> &'static str { + "prefix_hash" + } + + fn as_any(&self) -> &dyn std::any::Any { + self + } +} + +#[cfg(test)] +mod tests { + use super::*; + use crate::core::{BasicWorkerBuilder, HashRing, WorkerType}; + + fn create_workers(urls: &[&str]) -> Vec> { + urls.iter() + .map(|url| { + Arc::new( + BasicWorkerBuilder::new(*url) + .worker_type(WorkerType::Regular) + .build(), + ) as Arc + }) + .collect() + } + + #[test] + fn test_prefix_hash_consistent_routing() { + let policy = PrefixHashPolicy::with_defaults(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); + let ring = Arc::new(HashRing::new(&workers)); + + // Same tokens should always route to same worker + let tokens: Vec = vec![1, 2, 3, 4, 5, 6, 7, 8, 9, 10]; + let info = SelectWorkerInfo { + tokens: Some(&tokens), + hash_ring: Some(ring.clone()), + ..Default::default() + }; + + let (first_result, _) = policy.select_worker_impl(&workers, &info); + let first_idx = first_result.unwrap(); + + // Verify consistency + for _ in 0..10 { + let (result, _) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, Some(first_idx)); + } + } + + #[test] + fn test_different_prefixes_distribute() { + let policy = PrefixHashPolicy::with_defaults(); + let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); + let ring = Arc::new(HashRing::new(&workers)); + + let mut distribution = std::collections::HashMap::new(); + + // Different token sequences should distribute across workers + for i in 0..100 { + let tokens: Vec = vec![i, i + 1, i + 2, i + 3]; + let info = SelectWorkerInfo { + tokens: Some(&tokens), + hash_ring: Some(ring.clone()), + ..Default::default() + }; + + let (result, _) = policy.select_worker_impl(&workers, &info); + *distribution.entry(result.unwrap()).or_insert(0) += 1; + } + + assert!( + distribution.len() > 1, + "Should distribute across workers, got {:?}", + distribution + ); + } + + #[test] + fn test_shared_prefix_routes_same() { + let policy = PrefixHashPolicy::new(PrefixHashConfig { + prefix_token_count: 5, // Only look at first 5 tokens + ..Default::default() + }); + let workers = create_workers(&["http://w1:8000", "http://w2:8000", "http://w3:8000"]); + let ring = Arc::new(HashRing::new(&workers)); + + // Two sequences with same first 5 tokens should route to same worker + let tokens1: Vec = vec![1, 2, 3, 4, 5, 100, 200, 300]; + let tokens2: Vec = vec![1, 2, 3, 4, 5, 999, 888, 777]; + + let info1 = SelectWorkerInfo { + tokens: Some(&tokens1), + hash_ring: Some(ring.clone()), + ..Default::default() + }; + let info2 = SelectWorkerInfo { + tokens: Some(&tokens2), + hash_ring: Some(ring.clone()), + ..Default::default() + }; + + let (result1, _) = policy.select_worker_impl(&workers, &info1); + let (result2, _) = policy.select_worker_impl(&workers, &info2); + + assert_eq!(result1, result2, "Same prefix should route to same worker"); + } + + #[test] + fn test_no_tokens_returns_none() { + let policy = PrefixHashPolicy::with_defaults(); + let workers = create_workers(&["http://w1:8000"]); + let ring = Arc::new(HashRing::new(&workers)); + + // Empty tokens + let tokens: Vec = vec![]; + let info = SelectWorkerInfo { + tokens: Some(&tokens), + hash_ring: Some(ring.clone()), + ..Default::default() + }; + + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, None); + assert_eq!(branch, Branch::NoTokens); + + // No tokens field + let info_no_tokens = SelectWorkerInfo { + tokens: None, + hash_ring: Some(ring), + ..Default::default() + }; + + let (result2, branch2) = policy.select_worker_impl(&workers, &info_no_tokens); + assert_eq!(result2, None); + assert_eq!(branch2, Branch::NoTokens); + } + + #[test] + fn test_no_healthy_workers() { + let policy = PrefixHashPolicy::with_defaults(); + let workers = create_workers(&["http://w1:8000"]); + workers[0].set_healthy(false); + + let ring = Arc::new(HashRing::new(&workers)); + let tokens: Vec = vec![1, 2, 3]; + let info = SelectWorkerInfo { + tokens: Some(&tokens), + hash_ring: Some(ring), + ..Default::default() + }; + + let (result, branch) = policy.select_worker_impl(&workers, &info); + assert_eq!(result, None); + assert_eq!(branch, Branch::NoHealthyWorkers); + } + + #[test] + fn test_load_ok_calculation() { + let policy = PrefixHashPolicy::new(PrefixHashConfig { + load_factor: 1.25, + ..Default::default() + }); + + // Total load 100, 4 workers -> avg 25, threshold 31.25 + assert!(policy.load_ok(30, 100, 4)); // 30 <= 31.25 + assert!(!policy.load_ok(35, 100, 4)); // 35 > 31.25 + + // Edge cases + assert!(policy.load_ok(0, 0, 4)); // No load = OK + assert!(policy.load_ok(100, 0, 0)); // No workers = OK (shouldn't happen) + } + + #[test] + fn test_policy_name() { + let policy = PrefixHashPolicy::with_defaults(); + assert_eq!(policy.name(), "prefix_hash"); + } +} diff --git a/sgl-model-gateway/src/policies/registry.rs b/sgl-model-gateway/src/policies/registry.rs index 9e063282a..6b7c11354 100644 --- a/sgl-model-gateway/src/policies/registry.rs +++ b/sgl-model-gateway/src/policies/registry.rs @@ -9,10 +9,7 @@ use tracing::{debug, info, warn}; /// When the first worker of a new model is added, it determines the policy for that model. /// All subsequent workers of the same model use the established policy. /// When the last worker of a model is removed, the policy mapping is cleaned up. -use super::{ - BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy, - LoadBalancingPolicy, ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy, -}; +use super::{BucketPolicy, CacheAwarePolicy, LoadBalancingPolicy, PolicyFactory}; use crate::{config::types::PolicyConfig, core::Worker}; /// Registry for managing model-to-policy mappings @@ -160,60 +157,17 @@ impl PolicyRegistry { Arc::clone(&self.default_policy) } - /// Create a policy from a type string + /// Create a policy from a type string (delegates to PolicyFactory) fn create_policy_from_type(&self, policy_type: &str) -> Arc { - match policy_type { - "round_robin" => Arc::new(RoundRobinPolicy::new()), - "random" => Arc::new(RandomPolicy::new()), - "cache_aware" => Arc::new(CacheAwarePolicy::new()), - "power_of_two" => Arc::new(PowerOfTwoPolicy::new()), - "bucket" => Arc::new(BucketPolicy::new()), - "manual" => Arc::new(ManualPolicy::new()), - "consistent_hashing" => Arc::new(ConsistentHashingPolicy::new()), - _ => { - warn!("Unknown policy type '{}', using default", policy_type); - Arc::clone(&self.default_policy) - } - } + PolicyFactory::create_by_name(policy_type).unwrap_or_else(|| { + warn!("Unknown policy type '{}', using default", policy_type); + Arc::clone(&self.default_policy) + }) } - /// Create a policy from a PolicyConfig + /// Create a policy from a PolicyConfig (delegates to PolicyFactory) fn create_policy_from_config(config: &PolicyConfig) -> Arc { - match config { - PolicyConfig::RoundRobin => Arc::new(RoundRobinPolicy::new()), - PolicyConfig::Random => Arc::new(RandomPolicy::new()), - PolicyConfig::CacheAware { - cache_threshold, - balance_abs_threshold, - balance_rel_threshold, - eviction_interval_secs, - max_tree_size, - } => { - let cache_config = CacheAwareConfig { - cache_threshold: *cache_threshold, - balance_abs_threshold: *balance_abs_threshold, - balance_rel_threshold: *balance_rel_threshold, - eviction_interval_secs: *eviction_interval_secs, - max_tree_size: *max_tree_size, - }; - Arc::new(CacheAwarePolicy::with_config(cache_config)) - } - PolicyConfig::PowerOfTwo { .. } => Arc::new(PowerOfTwoPolicy::new()), - PolicyConfig::Bucket { - balance_abs_threshold, - balance_rel_threshold, - bucket_adjust_interval_secs, - } => { - let config = BucketConfig { - balance_abs_threshold: *balance_abs_threshold, - balance_rel_threshold: *balance_rel_threshold, - bucket_adjust_interval_secs: *bucket_adjust_interval_secs, - }; - Arc::new(BucketPolicy::with_config(config)) - } - PolicyConfig::Manual => Arc::new(ManualPolicy::new()), - PolicyConfig::ConsistentHashing => Arc::new(ConsistentHashingPolicy::new()), - } + PolicyFactory::create_from_config(config) } /// Get current model->policy mappings (for debugging/monitoring) diff --git a/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs b/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs index f19531543..d29780025 100644 --- a/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs +++ b/sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs @@ -67,11 +67,23 @@ impl PipelineStage for WorkerSelectionStage { prep.original_text.as_deref() }; + // Get tokens for PrefixHash policy support + let tokens = if prep.token_ids.is_empty() { + None + } else { + Some(prep.token_ids.as_slice()) + }; + let headers = ctx.input.headers.as_ref(); let workers = match self.mode { WorkerSelectionMode::Regular => { - match self.select_single_worker(ctx.input.model_id.as_deref(), text, headers) { + match self.select_single_worker( + ctx.input.model_id.as_deref(), + text, + tokens, + headers, + ) { Some(w) => WorkerSelection::Single { worker: w }, None => { error!( @@ -88,7 +100,7 @@ impl PipelineStage for WorkerSelectionStage { } } WorkerSelectionMode::PrefillDecode => { - match self.select_pd_pair(ctx.input.model_id.as_deref(), text, headers) { + match self.select_pd_pair(ctx.input.model_id.as_deref(), text, tokens, headers) { Some((prefill, decode)) => WorkerSelection::Dual { prefill, decode }, None => { error!( @@ -123,6 +135,7 @@ impl WorkerSelectionStage { &self, model_id: Option<&str>, text: Option<&str>, + tokens: Option<&[u32]>, headers: Option<&http::HeaderMap>, ) -> Option> { // Get workers for the specified model, filtered by connection mode @@ -158,6 +171,7 @@ impl WorkerSelectionStage { &available, &SelectWorkerInfo { request_text: text, + tokens, headers, hash_ring, }, @@ -179,6 +193,7 @@ impl WorkerSelectionStage { &self, model_id: Option<&str>, text: Option<&str>, + tokens: Option<&[u32]>, headers: Option<&http::HeaderMap>, ) -> Option<(Arc, Arc)> { let all_workers = self.worker_registry.get_workers_filtered( @@ -226,6 +241,7 @@ impl WorkerSelectionStage { let info = SelectWorkerInfo { request_text: text, + tokens, headers, hash_ring, }; diff --git a/sgl-model-gateway/src/routers/http/pd_router.rs b/sgl-model-gateway/src/routers/http/pd_router.rs index 8bceb261b..308f4eeee 100644 --- a/sgl-model-gateway/src/routers/http/pd_router.rs +++ b/sgl-model-gateway/src/routers/http/pd_router.rs @@ -806,6 +806,7 @@ impl PDRouter { &available_workers, &SelectWorkerInfo { request_text, + tokens: None, // HTTP doesn't have tokens, use gRPC for PrefixHash headers, hash_ring, }, diff --git a/sgl-model-gateway/src/routers/http/router.rs b/sgl-model-gateway/src/routers/http/router.rs index 8e61459ac..a8956e2a5 100644 --- a/sgl-model-gateway/src/routers/http/router.rs +++ b/sgl-model-gateway/src/routers/http/router.rs @@ -178,6 +178,7 @@ impl Router { &available, &SelectWorkerInfo { request_text: text, + tokens: None, // HTTP doesn't have tokens, use gRPC for PrefixHash headers, hash_ring, },