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

[HiCache][LoRA] Isolate storage pages by extra key (#38577)

Browse Source 
Co-authored-by: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com>
This commit is contained in:
Yanbin Jiang
2026-09-12 18:18:14 +08:00
committed by GitHub
co-authored by Shuwen Wang
parent b9cb96496d
commit 0b415fa573
21 changed files with 516 additions and 132 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/sglang-radix-tree/src/node.rs
+22 -2
View File
@@ -1137,7 +1137,7 @@ impl<K: ChildKeyType> NodeArena<K> {
.filter_map(|(idx, slot)| slot.as_ref().map(|_| NodeIdx_(idx)))
}
/// Per-page hash values for a node's key, chained from its parent's last hash.
/// Chain page hashes from the parent, or seed a new chain with the namespace.
pub fn compute_node_hash_values(&self, node_id: NodeIdx_, page_size: usize) -> Vec<String> {
let node = self.node(node_id);
let parent_hash = node.parent.and_then(|parent_id| {
@@ -1148,7 +1148,27 @@ impl<K: ChildKeyType> NodeArena<K> {
None
}
});
crate::node::get_hash_str::<K>(node.key.as_ref(), parent_hash, page_size)
let prior = parent_hash.map(str::to_owned).or_else(|| {
let namespace = node.namespace.as_ref();
if namespace == KeyNamespaceRef::default() {
return None;
}
// Match Python's storage_namespace_seed byte for byte.
let mut hasher = Sha256::new();
hasher.update(b"sglang-cache-namespace-v1");
for part in [namespace.extra_key, namespace.cache_salt] {
match part {
None => hasher.update([0u8]),
Some(part) => {
hasher.update([1u8]);
hasher.update((part.len() as u64).to_le_bytes());
hasher.update(part.as_bytes());
}
}
}
Some(digest_to_hex(&hasher.finalize().into()))
});
crate::node::get_hash_str::<K>(node.key.as_ref(), prior.as_deref(), page_size)
}
/// The ancestor chain's hash values ending at `node_id`, in root-to-node
rust/sglang-radix-tree/src/tests/node.rs
+21
View File
@@ -1945,3 +1945,24 @@ fn iter_yields_all_members() {
members.sort_unstable();
assert_eq!(members, vec![NodeIdx_(10), NodeIdx_(30)]);
}
// Pin Python/Rust storage hashes across a parent-child boundary.
#[test]
fn storage_hashes_match_python() -> Result<(), TreeCoreRuntimeError> {
let namespace = KeyNamespaceRef::new(Some("adapter-a"), Some("tenant-a"));
let mut arena: NodeArena<Vec<i64>> = NodeArena::new(vec![FULL], 2);
let root = arena.root();
let parent = arena.alloc_child_in_namespace(root, vec![1, 2], 0, namespace)?;
let hashes = arena.compute_node_hash_values(parent, 2);
assert_eq!(
hashes,
vec!["91b8b854063250a84c6f75b3d294bc5d72047c3a15c52b09038a5831f69ecd1a"]
);
arena.node_mut(parent).hash_value = Some(hashes);
let child = arena.alloc_child_in_namespace(parent, vec![3, 4], 0, namespace)?;
assert_eq!(
arena.compute_node_hash_values(child, 2),
vec!["c1ab67afa32b9fdd2ac8429d44d56f207a99c03a30b7a9bb131a32db35354c90"]
);
Ok(())
}
rust/sglang-radix-tree/src/tests/unified_tree_core.rs
+55 -14
View File
@@ -2754,11 +2754,11 @@ fn insert_coalesces_parent_linked_block_stores() {
.hash_value,
Some(hashes)
);
assert!(tc.salted_event_hashes.is_empty());
assert!(tc.namespaced_event_hashes.is_empty());
}
#[test]
fn salted_event_hashes_are_sparse_and_removed_with_the_node() {
fn namespaced_event_hashes_are_sparse_and_removed_with_the_node() {
let mut tc = events_core(2);
let key = vec![1, 2, 7, 8];
tc.insert(&insert_params_in_namespace(
@@ -2773,8 +2773,8 @@ fn salted_event_hashes_are_sparse_and_removed_with_the_node() {
.match_prefix(&match_params_in_namespace(&key, None, Some("tenant-a")))
.best_match_node_id;
let leaf_idx = tc.arena.resolve(leaf).expect("live test node");
assert_eq!(tc.salted_event_hashes[&leaf].len(), 2);
assert_eq!(
assert_eq!(tc.namespaced_event_hashes[&leaf].len(), 2);
assert_ne!(
tc.arena.node(leaf_idx).hash_value,
Some(crate::node::get_hash_str::<Vec<i64>>(&key, None, 2))
);
@@ -2790,7 +2790,7 @@ fn salted_event_hashes_are_sparse_and_removed_with_the_node() {
accumulate_step(step, &mut tracker, &mut device_frees, &mut host_frees);
tc.evict_device_end(FULL);
tc.take_events();
assert!(tc.salted_event_hashes.is_empty());
assert!(tc.namespaced_event_hashes.is_empty());
tc.insert(&insert_params_in_namespace(
&key,
@@ -2798,13 +2798,48 @@ fn salted_event_hashes_are_sparse_and_removed_with_the_node() {
None,
Some("tenant-a"),
));
assert!(!tc.salted_event_hashes.is_empty());
assert!(!tc.namespaced_event_hashes.is_empty());
tc.reset();
assert!(tc.salted_event_hashes.is_empty());
assert!(tc.namespaced_event_hashes.is_empty());
}
#[test]
fn salted_event_hashes_survive_node_split() {
fn extra_key_nodes_publish_token_only_event_hashes() {
// Events omit extra_key; storage includes it.
let mut tc = events_core(2);
let key = vec![1, 2, 7, 8];
tc.insert(&insert_params_in_namespace(
&key,
&[10, 11, 12, 13],
Some("lora-a"),
None,
));
let token_only = crate::node::get_hash_str::<Vec<i64>>(&key, None, 2);
assert_eq!(
tc.take_events(),
vec![KvCacheEvent::BlockStored {
block_hashes: token_only
.iter()
.map(|hash| crate::node::hash_str_to_int64(hash))
.collect(),
parent_block_hash: None,
token_ids: key.clone(),
block_size: 2,
medium: StorageMedium::Gpu,
cache_salt: None,
}]
);
let leaf = tc
.match_prefix(&match_params_in_namespace(&key, Some("lora-a"), None))
.best_match_node_id;
let leaf_idx = tc.arena.resolve(leaf).expect("live test node");
assert_eq!(tc.namespaced_event_hashes[&leaf].len(), 2);
assert_ne!(tc.arena.node(leaf_idx).hash_value, Some(token_only));
}
#[test]
fn namespaced_event_hashes_survive_node_split() {
let mut tc = events_core(2);
let original = vec![1, 2, 3, 4];
tc.insert(&insert_params_in_namespace(
@@ -2820,7 +2855,7 @@ fn salted_event_hashes_survive_node_split() {
Some("tenant-a"),
))
.best_match_node_id;
let original_hashes = tc.salted_event_hashes[&original_leaf].clone();
let original_hashes = tc.namespaced_event_hashes[&original_leaf].clone();
tc.take_events();
let branch = vec![1, 2, 5, 6];
@@ -2844,8 +2879,14 @@ fn salted_event_hashes_survive_node_split() {
.node(tc.arena.resolve(split_child).expect("live test node"))
.parent();
let split_parent = tc.arena.node(split_parent_idx).id;
assert_eq!(tc.salted_event_hashes[&split_parent], original_hashes[..1]);
assert_eq!(tc.salted_event_hashes[&split_child], original_hashes[1..]);
assert_eq!(
tc.namespaced_event_hashes[&split_parent],
original_hashes[..1]
);
assert_eq!(
tc.namespaced_event_hashes[&split_child],
original_hashes[1..]
);
}
#[test]
@@ -2863,9 +2904,9 @@ fn salted_event_hash_walk_is_iterative_and_on_demand() {
)
.unwrap();
}
assert!(tc.salted_event_hashes.is_empty());
tc.ensure_salted_event_hashes_(parent);
assert_eq!(tc.salted_event_hashes.len(), 1100);
assert!(tc.namespaced_event_hashes.is_empty());
tc.ensure_namespaced_event_hashes_(parent);
assert_eq!(tc.namespaced_event_hashes.len(), 1100);
}
#[test]
rust/sglang-radix-tree/src/unified_tree_core.rs
+44 -48
View File
@@ -542,9 +542,8 @@ pub struct UnifiedTreeCore<K: ChildKeyType> {
pub(crate) enable_kv_cache_events: bool,
/// Queued placement events, drained by take_events.
pub(crate) kv_event_queue: Vec<KvCacheEvent<K::Atom>>,
/// Namespace-aware event hashes, populated only for salted nodes whose
/// placement events are requested. Storage hashes remain on the nodes.
pub(crate) salted_event_hashes: HashMap<NodeId, Vec<HashDigest>>,
/// Namespaced event hashes, seeded only by cache_salt; events omit extra_key.
pub(crate) namespaced_event_hashes: HashMap<NodeId, Vec<HashDigest>>,
/// Hit count at which a node earns a host write-through backup.
pub(crate) write_through_threshold: i64,
@@ -723,7 +722,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
has_swa_host_pool: params.has_swa_host_pool,
enable_kv_cache_events: params.enable_kv_cache_events,
kv_event_queue: Vec::new(),
salted_event_hashes: HashMap::new(),
namespaced_event_hashes: HashMap::new(),
write_through_threshold: params.write_through_threshold,
swa_uuid_counter: 1,
device: params.device,
@@ -752,7 +751,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
self.write_back_coexist_reclaim_digest = 0;
self.lru_lists = Self::new_lru_lists();
self.full_evict_device_heap.clear();
self.salted_event_hashes.clear();
self.namespaced_event_hashes.clear();
self.ongoing_insert_walk_state = None;
}
@@ -1798,13 +1797,13 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
child.hash_value = child_hash;
self.arena.node_mut(new_node_id).hash_value = new_node_hash;
let child_handle = self.arena.node(child_id).id;
if let Some(mut parent_event_hashes) = self.salted_event_hashes.remove(&child_handle) {
if let Some(mut parent_event_hashes) = self.namespaced_event_hashes.remove(&child_handle) {
let child_event_hashes = parent_event_hashes.split_off(split_len / self.page_size);
parent_event_hashes.shrink_to_fit();
let new_node_handle = self.arena.node(new_node_id).id;
self.salted_event_hashes
self.namespaced_event_hashes
.insert(new_node_handle, parent_event_hashes);
self.salted_event_hashes
self.namespaced_event_hashes
.insert(child_handle, child_event_hashes);
}
@@ -2530,7 +2529,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
pub fn remove_leaf_from_parent_(&mut self, node_id: NodeIdx_) {
// Arena slots are reused, so discard tracking before freeing the node.
self.full_coexisting_host_nodes.discard(node_id);
self.salted_event_hashes
self.namespaced_event_hashes
.remove(&self.arena.node(node_id).id);
// The arena is the registry: freeing detaches by page key and recycles the slot.
self.arena
@@ -2767,18 +2766,13 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
}
}
/// Fill the salted external event-hash chain through `node_id`.
fn ensure_salted_event_hashes_(&mut self, node_id: NodeIdx_) {
let cache_salt = self
.arena
.node(node_id)
.namespace
.cache_salt_arc()
.expect("salted event hashing requires cache_salt");
let node_handle = self.arena.node(node_id).id;
if self.salted_event_hashes.contains_key(&node_handle) {
/// Fill the event chain through `node_id`, using only cache_salt.
fn ensure_namespaced_event_hashes_(&mut self, node_id: NodeIdx_) {
let node = self.arena.node(node_id);
if self.namespaced_event_hashes.contains_key(&node.id) {
return;
}
let namespace = node.namespace.clone();
let mut missing = Vec::new();
let mut cursor = Some(node_id);
@@ -2789,11 +2783,10 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
break;
}
assert_eq!(
node.namespace.cache_salt(),
Some(cache_salt.as_ref()),
"radix path contains mismatched cache_salt values"
node.namespace, namespace,
"radix path contains mismatched cache namespaces"
);
if let Some(hashes) = self.salted_event_hashes.get(&node.id) {
if let Some(hashes) = self.namespaced_event_hashes.get(&node.id) {
prior = hashes.last().copied();
break;
}
@@ -2801,12 +2794,14 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
cursor = node.try_parent();
}
let mut prior = prior.unwrap_or_else(|| {
let mut hasher = Sha256::new();
hasher.update(b"sglang-cache-salt-v1\0");
hasher.update(cache_salt.as_bytes());
hasher.finalize().into()
});
if prior.is_none() {
prior = namespace.cache_salt().map(|cache_salt| {
let mut hasher = Sha256::new();
hasher.update(b"sglang-cache-salt-v1\0");
hasher.update(cache_salt.as_bytes());
hasher.finalize().into()
});
}
for id in missing.into_iter().rev() {
let (handle, hashes) = {
let node = self.arena.node(id);
@@ -2814,15 +2809,15 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
node.id,
crate::node::get_hash_digests::<K>(
node.key.as_ref(),
Some(&prior),
prior.as_ref(),
self.page_size,
),
)
};
if let Some(last) = hashes.last() {
prior = *last;
prior = Some(*last);
}
self.salted_event_hashes.insert(handle, hashes);
self.namespaced_event_hashes.insert(handle, hashes);
}
}
@@ -2836,15 +2831,16 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
self.arena.node_mut(node_id).hash_value = Some(hash_values);
}
let cache_salt = self.arena.node(node_id).namespace.cache_salt_arc();
if cache_salt.is_some() {
self.ensure_salted_event_hashes_(node_id);
let namespaced = self.arena.node(node_id).namespace != KeyNamespace::default();
if namespaced {
self.ensure_namespaced_event_hashes_(node_id);
}
let events = {
let node = self.arena.node(node_id);
let mut parent_block_hash = node.parent.and_then(|parent_id| {
let parent = self.arena.node(parent_id);
if cache_salt.is_some() {
self.salted_event_hashes
if namespaced {
self.namespaced_event_hashes
.get(&parent.id)
.and_then(|hashes| hashes.last())
.map(crate::node::hash_digest_to_int64)
@@ -2867,8 +2863,8 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
});
parent_block_hash = Some(block_hash);
};
if cache_salt.is_some() {
let hashes = &self.salted_event_hashes[&node.id];
if namespaced {
let hashes = &self.namespaced_event_hashes[&node.id];
assert!(
hashes.len() >= num_pages,
"store event: {} page hashes for {num_pages} pages",
@@ -2904,14 +2900,14 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
let hash_values = self.arena.compute_node_hash_values(node_id, self.page_size);
self.arena.node_mut(node_id).hash_value = Some(hash_values);
}
let cache_salt = self.arena.node(node_id).namespace.cache_salt_arc();
if cache_salt.is_some() {
self.ensure_salted_event_hashes_(node_id);
let namespaced = self.arena.node(node_id).namespace != KeyNamespace::default();
if namespaced {
self.ensure_namespaced_event_hashes_(node_id);
}
let node = self.arena.node(node_id);
let num_pages = node.key.atom_len().div_ceil(self.page_size);
let block_hashes: Vec<i64> = if cache_salt.is_some() {
self.salted_event_hashes[&node.id][..num_pages]
let block_hashes: Vec<i64> = if namespaced {
self.namespaced_event_hashes[&node.id][..num_pages]
.iter()
.map(crate::node::hash_digest_to_int64)
.collect()
@@ -3955,23 +3951,23 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
&orphans[..orphans.len().min(5)]
));
}
for (&node_handle, hashes) in &self.salted_event_hashes {
for (&node_handle, hashes) in &self.namespaced_event_hashes {
let Ok(node_id) = self.arena.resolve(node_handle) else {
errors.push(format!(
"[Events] salted hashes reference freed node {node_handle}"
"[Events] event hashes reference freed node {node_handle}"
));
continue;
};
let node = self.arena.node(node_id);
if node.namespace.cache_salt().is_none() {
if node.namespace == KeyNamespace::default() {
errors.push(format!(
"[Events] unsalted node {node_handle} carries salted hashes"
"[Events] default-namespace node {node_handle} carries event hashes"
));
}
let expected_pages = node.key.atom_len().div_ceil(self.page_size);
if hashes.len() != expected_pages {
errors.push(format!(
"[Events] node {node_handle} has {} salted hashes for {expected_pages} pages",
"[Events] node {node_handle} has {} event hashes for {expected_pages} pages",
hashes.len()
));
}