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

[model-gateway] Optimize INSERT with leaf-only timestamp updates (#16097)

Browse Source
This commit is contained in:
Simo Lin
2025-12-29 10:43:20 -08:00
committed by GitHub
parent f44c4b3715
commit 684e148e38
1 changed files with 31 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files
sgl-model-gateway/src/policies/tree.rs
+31 -44
View File
@@ -302,18 +302,14 @@ impl Tree {
// Insert text into tree with given tenant // Insert text into tree with given tenant
// Use slice-based traversal to avoid Vec<char> allocation // Use slice-based traversal to avoid Vec<char> allocation
// Get epoch for LRU ordering
let epoch = get_epoch();
// Intern the tenant ID once for reuse // Intern the tenant ID once for reuse
let tenant_id = intern_tenant(tenant); let tenant_id = intern_tenant(tenant);
// Pre-compute total char count once (needed for size tracking) // Ensure tenant exists at root (don't update timestamp - root is never evicted)
let total_char_count = text.chars().count();
self.root self.root
.tenant_last_access_time .tenant_last_access_time
.insert(Arc::clone(&tenant_id), epoch); .entry(Arc::clone(&tenant_id))
.or_insert(0);
self.tenant_char_count self.tenant_char_count
.entry(Arc::clone(&tenant_id)) .entry(Arc::clone(&tenant_id))
@@ -321,7 +317,6 @@ impl Tree {
// Track remaining text as a slice - no allocation needed // Track remaining text as a slice - no allocation needed
let mut remaining = text; let mut remaining = text;
let mut matched_chars = 0;
let mut prev = Arc::clone(&self.root); let mut prev = Arc::clone(&self.root);
// Result type to carry state out of the match block // Result type to carry state out of the match block
@@ -340,8 +335,11 @@ impl Tree {
// Use entry API for atomic check-and-insert semantics (required for thread safety) // Use entry API for atomic check-and-insert semantics (required for thread safety)
let step = match prev.children.entry(first_char) { let step = match prev.children.entry(first_char) {
Entry::Vacant(entry) => { Entry::Vacant(entry) => {
// No match - create new node with remaining text // No match - create new node with remaining text (this is the leaf)
let remaining_char_count = total_char_count - matched_chars; // Compute remaining char count lazily - only here when creating leaf
let remaining_char_count = remaining.chars().count();
let epoch = get_epoch();
let new_node = Arc::new(Node { let new_node = Arc::new(Node {
children: DashMap::with_hasher(CharHasherBuilder::default()), children: DashMap::with_hasher(CharHasherBuilder::default()),
text: RwLock::new(NodeText::new(remaining.to_string())), text: RwLock::new(NodeText::new(remaining.to_string())),
@@ -349,7 +347,7 @@ impl Tree {
parent: RwLock::new(Some(Arc::clone(&prev))), parent: RwLock::new(Some(Arc::clone(&prev))),
}); });
// Attach tenant to the new node and increment count // Attach tenant to the new leaf node with timestamp
self.tenant_char_count self.tenant_char_count
.entry(Arc::clone(&tenant_id)) .entry(Arc::clone(&tenant_id))
.and_modify(|count| *count += remaining_char_count) .and_modify(|count| *count += remaining_char_count)
@@ -398,21 +396,14 @@ impl Tree {
*matched_node.text.write().unwrap() = contracted_text; *matched_node.text.write().unwrap() = contracted_text;
*matched_node.parent.write().unwrap() = Some(Arc::clone(&new_node)); *matched_node.parent.write().unwrap() = Some(Arc::clone(&new_node));
// Atomically attach tenant to the new split node // Attach tenant to the new split node (intermediate - no timestamp update)
match new_node // The cloned DashMap already has the tenant; just ensure char count is correct
.tenant_last_access_time if !new_node.tenant_last_access_time.contains_key(tenant_id.as_ref()) {
.entry(Arc::clone(&tenant_id)) self.tenant_char_count
{ .entry(Arc::clone(&tenant_id))
Entry::Vacant(v) => { .and_modify(|count| *count += matched_text_count)
self.tenant_char_count .or_insert(matched_text_count);
.entry(Arc::clone(&tenant_id)) new_node.tenant_last_access_time.insert(Arc::clone(&tenant_id), 0);
.and_modify(|count| *count += matched_text_count)
.or_insert(matched_text_count);
v.insert(epoch);
}
Entry::Occupied(mut o) => {
o.insert(epoch);
}
} }
InsertStep::Continue { InsertStep::Continue {
@@ -420,24 +411,16 @@ impl Tree {
advance_chars: shared_count, advance_chars: shared_count,
} }
} else { } else {
// Full match - move to next node // Full match - move to next node (intermediate - no timestamp update)
drop(matched_node_text); drop(matched_node_text);
// Atomically attach tenant to existing node // Ensure tenant exists at this intermediate node
match matched_node if !matched_node.tenant_last_access_time.contains_key(tenant_id.as_ref()) {
.tenant_last_access_time self.tenant_char_count
.entry(Arc::clone(&tenant_id)) .entry(Arc::clone(&tenant_id))
{ .and_modify(|count| *count += matched_node_text_count)
Entry::Vacant(v) => { .or_insert(matched_node_text_count);
self.tenant_char_count matched_node.tenant_last_access_time.insert(Arc::clone(&tenant_id), 0);
.entry(Arc::clone(&tenant_id))
.and_modify(|count| *count += matched_node_text_count)
.or_insert(matched_node_text_count);
v.insert(epoch);
}
Entry::Occupied(mut o) => {
o.insert(epoch);
}
} }
InsertStep::Continue { InsertStep::Continue {
@@ -450,17 +433,21 @@ impl Tree {
// Entry guard is now dropped - safe to update prev // Entry guard is now dropped - safe to update prev
match step { match step {
InsertStep::Done => break, InsertStep::Done => return, // New leaf created with timestamp, we're done
InsertStep::Continue { InsertStep::Continue {
next_prev, next_prev,
advance_chars, advance_chars,
} => { } => {
prev = next_prev; prev = next_prev;
remaining = advance_by_chars(remaining, advance_chars); remaining = advance_by_chars(remaining, advance_chars);
matched_chars += advance_chars;
} }
} }
} }
// Loop exited normally (remaining empty) - prev is the leaf node
// Update its timestamp for LRU ordering
let epoch = get_epoch();
prev.tenant_last_access_time.insert(Arc::clone(&tenant_id), epoch);
} }
/// Performs prefix matching and returns detailed result with char counts. /// Performs prefix matching and returns detailed result with char counts.