chenchenghao/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

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

Browse Source
This commit is contained in:
Simo Lin
2025-12-27 05:58:05 -05:00
committed by GitHub
parent ca740a41f3
commit 3645ed0f73
14 changed files with 532 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
sgl-model-gateway/Cargo.toml
View File

@@ -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"

1
sgl-model-gateway/bindings/python/sglang_router/router.py
View File

@@ -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]

5
sgl-model-gateway/bindings/python/src/lib.rs
View File

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

25
sgl-model-gateway/src/config/types.rs
View File

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

20
sgl-model-gateway/src/config/validation.rs
View File

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

18
sgl-model-gateway/src/main.rs
View File

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

9
sgl-model-gateway/src/observability/metrics.rs
View File

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

14
sgl-model-gateway/src/policies/factory.rs
View File

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

5
sgl-model-gateway/src/policies/mod.rs
View File

@@ -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

409
sgl-model-gateway/src/policies/prefix_hash.rs Normal file
View File

@@ -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");
}
}

62
sgl-model-gateway/src/policies/registry.rs
View File

@@ -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<dyn LoadBalancingPolicy> {
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<dyn LoadBalancingPolicy> {
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)

20
sgl-model-gateway/src/routers/grpc/common/stages/worker_selection.rs
View File

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

1
sgl-model-gateway/src/routers/http/pd_router.rs
View File

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

1
sgl-model-gateway/src/routers/http/router.rs
View File

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