[model-gateway] Add PrefixHash load balancing policy for KV cache-aware routing (#15935)

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"

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]
 

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,
+                },
             }
         };
 

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",
         }
     }
 }

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(())
     }

sgl-model-gateway/src/main.rs

@@ -137,7 +137,7 @@ struct CliArgs {
     #[arg(long, num_args = 0..)]
     worker_urls: Vec<String>,
 
-    #[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<String>,
 
-    #[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<String>,
 
-    #[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<String>,
 
     #[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,
         }

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!(

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

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

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<dyn Worker>],
+        info: &SelectWorkerInfo,
+        prefix_hash: u64,
+    ) -> (Option<usize>, Branch) {
+        // Build healthy worker URL to index map
+        let healthy_workers: Vec<(usize, &Arc<dyn Worker>)> = 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<dyn Worker>)> =
+                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<dyn Worker>],
+        info: &SelectWorkerInfo,
+    ) -> (Option<usize>, 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<dyn Worker>], info: &SelectWorkerInfo) -> Option<usize> {
+        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<Arc<dyn Worker>> {
+        urls.iter()
+            .map(|url| {
+                Arc::new(
+                    BasicWorkerBuilder::new(*url)
+                        .worker_type(WorkerType::Regular)
+                        .build(),
+                ) as Arc<dyn Worker>
+            })
+            .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<u32> = 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<u32> = 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<u32> = vec![1, 2, 3, 4, 5, 100, 200, 300];
+        let tokens2: Vec<u32> = 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<u32> = 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<u32> = 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");
+    }
+}

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::{
+use super::{BucketPolicy, CacheAwarePolicy, LoadBalancingPolicy, PolicyFactory};
-    BucketConfig, BucketPolicy, CacheAwareConfig, CacheAwarePolicy, ConsistentHashingPolicy,
-    LoadBalancingPolicy, ManualPolicy, PowerOfTwoPolicy, RandomPolicy, RoundRobinPolicy,
-};
 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<dyn LoadBalancingPolicy> {
-        match policy_type {
+        PolicyFactory::create_by_name(policy_type).unwrap_or_else(|| {
-            "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)
-            }
+        })
-        }
     }
 
-    /// Create a policy from a PolicyConfig
+    /// Create a policy from a PolicyConfig (delegates to PolicyFactory)
     fn create_policy_from_config(config: &PolicyConfig) -> Arc<dyn LoadBalancingPolicy> {
-        match config {
+        PolicyFactory::create_from_config(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()),
-        }
     }
 
     /// Get current model->policy mappings (for debugging/monitoring)

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<Arc<dyn Worker>> {
         // 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<dyn Worker>, Arc<dyn Worker>)> {
         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,
         };

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

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