From 9940c6f592ad5773d9a0d4c88b83fea6e71c1189 Mon Sep 17 00:00:00 2001 From: Simo Lin Date: Tue, 30 Dec 2025 12:28:01 -0800 Subject: [PATCH] [model-gateway]: ASCII byte comparison and probabilistic timestamp updates (#16181) --- sgl-model-gateway/src/policies/tree.rs | 82 ++++++++++++++++++++++---- 1 file changed, 72 insertions(+), 10 deletions(-) diff --git a/sgl-model-gateway/src/policies/tree.rs b/sgl-model-gateway/src/policies/tree.rs index fd178b0d1..4215dfbd5 100644 --- a/sgl-model-gateway/src/policies/tree.rs +++ b/sgl-model-gateway/src/policies/tree.rs @@ -13,6 +13,28 @@ use tracing::debug; type NodeRef = Arc; +/// Shard counts for DashMaps to balance concurrency vs allocation overhead. +/// Default DashMap uses num_cpus * 4 shards (e.g., 256 on 64-core machines). +/// +/// Root node uses higher shard count since ALL requests pass through it. +/// Other nodes use lower count as traffic diverges through the tree. +/// +/// This reduces memory by ~90% vs default while maintaining good concurrency. +const ROOT_SHARD_COUNT: usize = 32; +const NODE_SHARD_COUNT: usize = 8; + +/// Create a children DashMap for non-root nodes +#[inline] +fn new_children_map() -> DashMap { + DashMap::with_hasher_and_shard_amount(CharHasherBuilder::default(), NODE_SHARD_COUNT) +} + +/// Create a tenant access time DashMap for non-root nodes +#[inline] +fn new_tenant_map() -> DashMap { + DashMap::with_shard_amount(NODE_SHARD_COUNT) +} + /// Interned tenant ID to avoid repeated string allocations. /// Using Arc allows cheap cloning and comparison. pub type TenantId = Arc; @@ -65,11 +87,22 @@ type CharHasherBuilder = BuildHasherDefault; /// Advance a string slice by N characters, returning the remaining slice. /// Returns empty string if n >= char count. +/// Optimized: uses direct byte slicing for ASCII, falls back to char_indices for UTF-8. #[inline] fn advance_by_chars(s: &str, n: usize) -> &str { if n == 0 { return s; } + if n >= s.len() { + return ""; + } + // Fast path: if first N bytes are all ASCII, we can slice directly + let bytes = s.as_bytes(); + if bytes[..n].is_ascii() { + // Safe: we verified all bytes in [0..n] are ASCII (valid UTF-8 boundary) + return &s[n..]; + } + // Slow path: UTF-8 requires char-by-char traversal s.char_indices() .nth(n) .map(|(idx, _)| &s[idx..]) @@ -251,9 +284,31 @@ impl PartialEq for EvictionEntry { /// Count matching prefix characters between two strings. /// Returns the number of characters that match from the start. -/// Uses iterator-based comparison - no allocation required. +/// Optimized: uses fast byte comparison for ASCII, falls back to char iteration for UTF-8. #[inline] fn shared_prefix_count(a: &str, b: &str) -> usize { + let a_bytes = a.as_bytes(); + let b_bytes = b.as_bytes(); + + // Find common byte prefix length using iterator (potentially SIMD-optimized) + let common_byte_len = a_bytes + .iter() + .zip(b_bytes) + .position(|(&a_byte, &b_byte)| a_byte != b_byte) + .unwrap_or_else(|| a_bytes.len().min(b_bytes.len())); + + // If the common byte prefix is all ASCII, byte count == char count + // Otherwise, fall back to char-by-char comparison for UTF-8 safety + if a_bytes[..common_byte_len].is_ascii() { + common_byte_len + } else { + shared_prefix_count_chars(a, b) + } +} + +/// Fallback char-by-char comparison for strings with non-ASCII characters. +#[inline] +fn shared_prefix_count_chars(a: &str, b: &str) -> usize { a.chars() .zip(b.chars()) .take_while(|(a_char, b_char)| a_char == b_char) @@ -289,14 +344,18 @@ impl Tree { pub fn new() -> Self { Tree { + // Root uses higher shard count since ALL requests pass through it root: Arc::new(Node { - children: DashMap::with_hasher(CharHasherBuilder::default()), + children: DashMap::with_hasher_and_shard_amount( + CharHasherBuilder::default(), + ROOT_SHARD_COUNT, + ), text: RwLock::new(NodeText::empty()), - tenant_last_access_time: DashMap::new(), + tenant_last_access_time: DashMap::with_shard_amount(ROOT_SHARD_COUNT), parent: RwLock::new(None), last_tenant: parking_lot::RwLock::new(None), }), - tenant_char_count: DashMap::new(), + tenant_char_count: DashMap::with_shard_amount(ROOT_SHARD_COUNT), } } @@ -343,9 +402,9 @@ impl Tree { let epoch = get_epoch(); let new_node = Arc::new(Node { - children: DashMap::with_hasher(CharHasherBuilder::default()), + children: new_children_map(), text: RwLock::new(NodeText::new(remaining.to_string())), - tenant_last_access_time: DashMap::new(), + tenant_last_access_time: new_tenant_map(), parent: RwLock::new(Some(Arc::clone(&prev))), last_tenant: parking_lot::RwLock::new(Some(Arc::clone(&tenant_id))), }); @@ -384,7 +443,7 @@ impl Tree { let new_node = Arc::new(Node { text: RwLock::new(matched_text), - children: DashMap::with_hasher(CharHasherBuilder::default()), + children: new_children_map(), parent: RwLock::new(Some(Arc::clone(&prev))), tenant_last_access_time: matched_node.tenant_last_access_time.clone(), last_tenant: parking_lot::RwLock::new( @@ -539,10 +598,13 @@ impl Tree { } }; - // Update timestamp on the matched node only (O(1)). + // Update timestamp probabilistically (1 in 8 matches) to reduce DashMap contention. + // LRU eviction doesn't need perfect accuracy - approximate timestamps suffice. let epoch = get_epoch(); - curr.tenant_last_access_time - .insert(Arc::clone(&tenant), epoch); + if epoch & 0x7 == 0 { + curr.tenant_last_access_time + .insert(Arc::clone(&tenant), epoch); + } // Compute input char count from matched + remaining (deferred from start) let input_char_count = matched_chars + remaining.chars().count();