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

[Unified Tree] Port SWA Branching-Point Caching to the Rust TreeCore (#37584)

Browse Source
This commit is contained in:
Shuwen Wang
2026-09-12 18:37:57 +08:00
committed by GitHub
parent 0b415fa573
commit bd45cd50ca
15 changed files with 762 additions and 109 deletions
Download Patch File Download Diff File Expand all files Collapse all files
rust/sglang-radix-tree/src/components/swa.rs
+127 -14
View File
@@ -152,6 +152,60 @@ impl SwaComponent {
}
}
/// Nodes whose SWA data needs a host backup, deepest first. Buffer mode
/// stages one node per FIFO backup intent; cache mode backs up every
/// device-only node within one sliding window of `node_id`.
fn collect_unbacked_swa_nodes_<K: ChildKeyType>(
&self,
tree_core: &UnifiedTreeCore<K>,
node_id: NodeIdx_,
) -> Vec<NodeIdx_> {
if !tree_core.has_swa_host_pool {
return Vec::new();
}
if tree_core.is_host_memory_buffer_only {
return if tree_core.arena.node(node_id).has_device_value(SWA) {
vec![node_id]
} else {
Vec::new()
};
}
self.collect_unbacked_swa_nodes_in_window_(tree_core, node_id)
}
/// Nodes within one sliding window of `node_id` whose SWA data sits on
/// device with no host copy, deepest first. The walk stops at a node an
/// in-flight backup already covers: that ack owns everything above it, so
/// two acks can never claim the same node.
fn collect_unbacked_swa_nodes_in_window_<K: ChildKeyType>(
&self,
tree_core: &UnifiedTreeCore<K>,
node_id: NodeIdx_,
) -> Vec<NodeIdx_> {
let mut covered_tokens = 0;
let mut unbacked: Vec<NodeIdx_> = Vec::new();
let mut cur_id = node_id;
while covered_tokens < self.sliding_window_size {
let cur = tree_core.arena.node(cur_id);
if cur.is_root() || cur.write_through_pending_id.is_some() {
break;
}
let (on_device, on_host) = (cur.has_device_value(SWA), cur.has_host_value(SWA));
covered_tokens += if on_device {
cur.device_value_len(SWA)
} else if on_host {
cur.host_value_len(SWA)
} else {
break;
};
if on_device && !on_host {
unbacked.push(cur_id);
}
cur_id = cur.parent();
}
unbacked
}
fn next_host_unlocked_device_lru_node<K: ChildKeyType>(
tree_core: &UnifiedTreeCore<K>,
from: Option<NodeIdx_>,
@@ -305,6 +359,12 @@ impl SwaComponent {
}
impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
fn needs_incremental_backup(&self, tree_core: &UnifiedTreeCore<K>, node_id: NodeIdx_) -> bool {
!self
.collect_unbacked_swa_nodes_(tree_core, node_id)
.is_empty()
}
fn component_type(&self) -> ComponentType {
SWA
}
@@ -363,6 +423,14 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
value_chunks: &[Tensor],
best_value_len: usize,
) -> MatchResult {
let swa_boundary_len = result.device_indices.size()[0] as usize + result.host_hit_length;
// Branch at the last page-aligned Full-KV position past the SWA boundary.
let page_aligned_full_hit_len =
result.full_kv_hit_length / tree_core.page_size * tree_core.page_size;
result.swa_branching_seqlen =
(page_aligned_full_hit_len > swa_boundary_len).then_some(page_aligned_full_hit_len);
// Sum the SWA tokens backing the match, walking up from the best match
// until one sliding window is covered; host-resident chunks count
// toward the SWA host hit.
@@ -536,6 +604,10 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
result: &mut InsertResult,
cache_actions: &mut Vec<CacheAction>,
) {
if let Some(branching_seqlen) = params.swa_branching_seqlen {
result.swa_branch_inserted = params.key.atom_len() >= branching_seqlen;
}
if !is_new_leaf {
return;
}
@@ -864,19 +936,29 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
}
Ok(match phase {
CacheTransferPhase::BackupHost => {
let node = tree_core.arena.node(node_id);
if node.has_host_value(SWA) {
let unbacked_swa_nodes = self.collect_unbacked_swa_nodes_(tree_core, node_id);
if unbacked_swa_nodes.is_empty() {
return Ok(None);
}
// cd.value already holds SWA-pool indices (translated at insert time).
// Host pool indexing wants int64.
node.try_device_value(SWA).map(|value| {
vec![PoolTransfer {
name: PoolName::Swa,
device_indices: Some(value.to_kind(Kind::Int64)),
..Default::default()
}]
})
// Ancestors first: the host span is contiguous and the commit
// scatters it back in this order. Device values already hold
// SWA-pool indices (translated at insert time); host pool
// indexing wants int64.
let (device_indices, backup_node_ids): (Vec<Tensor>, Vec<NodeId>) =
unbacked_swa_nodes
.iter()
.rev()
.map(|&idx| {
let node = tree_core.arena.node(idx);
(node.device_value(SWA).to_kind(Kind::Int64), node.id)
})
.unzip();
Some(vec![PoolTransfer {
name: PoolName::Swa,
device_indices: Some(Tensor::cat(&device_indices, 0)),
nodes_to_load: Some(backup_node_ids),
..Default::default()
}])
}
CacheTransferPhase::LoadBack => {
// `node` is best_match_node; the SWA validator guarantees every
@@ -996,14 +1078,45 @@ impl<K: ChildKeyType> TreeComponent<K> for SwaComponent {
) {
match phase {
CacheTransferPhase::BackupHost => {
if let Some(transfer) = transfers.first()
&& let Some(host_indices) = &transfer.host_indices
{
let Some(transfer) = transfers.first() else {
return;
};
let Some(host_indices) = &transfer.host_indices else {
return;
};
// A missing or empty `nodes_to_load` means the span is this node's alone.
let target_ids = transfer
.nodes_to_load
.as_deref()
.filter(|ids| !ids.is_empty());
let Some(target_ids) = target_ids else {
let node = tree_core.arena.node_mut(node_id);
if !node.has_host_value(SWA) {
node.set_host_value(SWA, host_indices.copy());
}
return;
};
let mut offset = 0i64;
for &target_id in target_ids {
let target_idx = tree_core
.arena
.resolve(target_id)
.expect("backup transfers must reference live nodes");
let target = tree_core.arena.node(target_idx);
assert!(
target.has_device_value(SWA) && !target.has_host_value(SWA),
"SWA backup target {} is not device-only",
target.id
);
let size = target.device_value_len(SWA) as i64;
tree_core.arena.set_host_value(
target_idx,
SWA,
host_indices.narrow(0, offset, size).copy(),
);
offset += size;
}
assert_eq!(offset, host_indices.size()[0]);
}
CacheTransferPhase::LoadBack => {
let transfer = transfers
rust/sglang-radix-tree/src/python_bindings.rs
+34 -1
View File
@@ -301,6 +301,8 @@ struct InspectionMatchResultInput {
#[pyo3(attribute)]
swa_host_hit_length: usize,
#[pyo3(attribute)]
swa_branching_seqlen: Option<usize>,
#[pyo3(attribute)]
mamba_host_hit_length: usize,
#[pyo3(attribute)]
mamba_branching_seqlen: Option<usize>,
@@ -515,6 +517,7 @@ pub struct InsertParamsBinding {
pub mamba_value: Option<Py<PyAny>>,
pub prev_prefix_len: usize,
pub swa_evicted_seqlen: usize,
pub swa_branching_seqlen: Option<usize>,
pub chunked: bool,
pub priority: i64,
pub track_adopted_ranges: bool,
@@ -523,7 +526,7 @@ pub struct InsertParamsBinding {
#[pymethods]
impl InsertParamsBinding {
#[new]
#[pyo3(signature = (key, value, extra_key = None, cache_salt = None, prev_prefix_len = 0, swa_evicted_seqlen = 0, chunked = false, priority = 0, mamba_value = None, track_adopted_ranges = false))]
#[pyo3(signature = (key, value, extra_key = None, cache_salt = None, prev_prefix_len = 0, swa_evicted_seqlen = 0, swa_branching_seqlen = None, chunked = false, priority = 0, mamba_value = None, track_adopted_ranges = false))]
fn new(
py: Python<'_>,
key: &Bound<'_, PyAny>,
@@ -532,6 +535,7 @@ impl InsertParamsBinding {
cache_salt: Option<String>,
prev_prefix_len: usize,
swa_evicted_seqlen: usize,
swa_branching_seqlen: Option<usize>,
chunked: bool,
priority: i64,
mamba_value: Option<Py<PyAny>>,
@@ -545,6 +549,7 @@ impl InsertParamsBinding {
mamba_value,
prev_prefix_len,
swa_evicted_seqlen,
swa_branching_seqlen,
chunked,
priority,
track_adopted_ranges,
@@ -561,6 +566,7 @@ pub struct MatchResultBinding {
best_match_node_id: NodeId,
host_hit_length: usize,
swa_host_hit_length: usize,
swa_branching_seqlen: Option<usize>,
mamba_host_hit_length: usize,
mamba_branching_seqlen: Option<usize>,
full_kv_hit_length: usize,
@@ -577,6 +583,7 @@ impl MatchResultBinding {
best_match_node_id: result.best_match_node_id,
host_hit_length: result.host_hit_length,
swa_host_hit_length: result.swa_host_hit_length,
swa_branching_seqlen: result.swa_branching_seqlen,
mamba_host_hit_length: result.mamba_host_hit_length,
mamba_branching_seqlen: result.mamba_branching_seqlen,
full_kv_hit_length: result.full_kv_hit_length,
@@ -594,6 +601,7 @@ pub struct InsertResultBinding {
inserted_host_node: Option<NodeId>,
host_insert_dropped: bool,
mamba_exist: bool,
swa_branch_inserted: bool,
adopted_ranges: Option<HashMap<u8, Vec<(usize, usize)>>>,
cache_actions: Py<PyList>,
}
@@ -633,6 +641,7 @@ impl InsertResultBinding {
inserted_host_node: result.inserted_host_node,
host_insert_dropped: result.host_insert_dropped,
mamba_exist: result.mamba_exist,
swa_branch_inserted: result.swa_branch_inserted,
adopted_ranges: result.adopted_ranges.map(|ranges| {
ranges
.into_iter()
@@ -1018,6 +1027,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
mamba_value,
prev_prefix_len: params.prev_prefix_len,
swa_evicted_seqlen: params.swa_evicted_seqlen,
swa_branching_seqlen: params.swa_branching_seqlen,
chunked: params.chunked,
priority: params.priority,
track_adopted_ranges: params.track_adopted_ranges,
@@ -1053,6 +1063,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
mamba_value,
prev_prefix_len: params.prev_prefix_len,
swa_evicted_seqlen: params.swa_evicted_seqlen,
swa_branching_seqlen: params.swa_branching_seqlen,
chunked: params.chunked,
priority: params.priority,
track_adopted_ranges: params.track_adopted_ranges,
@@ -1355,6 +1366,16 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
.map_err(node_access_error)
}
/// Mark the host tier as buffer-only; wired after the host pools are built.
fn set_host_memory_buffer_only(&self, py: Python<'_>) {
py.allow_threads(|| self.core().set_host_memory_buffer_only());
}
/// Whether the host tier runs as a storage staging buffer, not a cache.
fn is_host_memory_buffer_only(&self, py: Python<'_>) -> bool {
py.allow_threads(|| self.core().is_host_memory_buffer_only)
}
/// Mark the host tier (HiCache) as wired.
fn set_hicache_enabled(&self, py: Python<'_>) {
py.allow_threads(|| self.core().set_hicache_enabled());
@@ -2300,6 +2321,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
best_match_node: best_match_node_id,
host_hit_length,
swa_host_hit_length,
swa_branching_seqlen,
mamba_host_hit_length,
mamba_branching_seqlen,
full_kv_hit_length,
@@ -2311,6 +2333,7 @@ impl<K: ChildKeyType + Send + Sync> TreeCoreBinding<K> {
best_match_node_id,
host_hit_length,
swa_host_hit_length,
swa_branching_seqlen,
mamba_host_hit_length,
mamba_branching_seqlen,
full_kv_hit_length,
@@ -2613,6 +2636,16 @@ macro_rules! tree_core_binding {
self.inner.is_full_device_evicted(py, node_id)
}
/// Mark the host tier as buffer-only; wired after the host pools are built.
fn set_host_memory_buffer_only(&self, py: Python<'_>) {
self.inner.set_host_memory_buffer_only(py)
}
/// Whether the host tier runs as a storage staging buffer, not a cache.
fn is_host_memory_buffer_only(&self, py: Python<'_>) -> bool {
self.inner.is_host_memory_buffer_only(py)
}
/// Mark the host tier (HiCache) as wired.
fn set_hicache_enabled(&self, py: Python<'_>) {
self.inner.set_hicache_enabled(py)
rust/sglang-radix-tree/src/tests/components/base.rs
+1
View File
@@ -95,6 +95,7 @@ fn insert_overlap_default_consumes_nothing() {
mamba_value: None,
prev_prefix_len: 0,
swa_evicted_seqlen: 0,
swa_branching_seqlen: None,
chunked: false,
priority: 0,
track_adopted_ranges: false,
rust/sglang-radix-tree/src/tests/components/full.rs
+2
View File
@@ -72,6 +72,7 @@ fn insert(tc: &mut UnifiedTreeCore<Vec<i64>>, key: &Vec<i64>, value: &[i64]) {
mamba_value: None,
prev_prefix_len: 0,
swa_evicted_seqlen: 0,
swa_branching_seqlen: None,
chunked: false,
priority: 0,
track_adopted_ranges: false,
@@ -473,6 +474,7 @@ fn host_drive_is_a_noop_without_host_leaves() {
mamba_value: None,
prev_prefix_len: 0,
swa_evicted_seqlen: 0,
swa_branching_seqlen: None,
chunked: false,
priority: 0,
track_adopted_ranges: false,
rust/sglang-radix-tree/src/tests/components/mamba.rs
+21
View File
@@ -65,6 +65,7 @@ fn insert_params_mamba<'k>(
mamba_value: mamba_slot.map(|slot| Tensor::from_slice(&[slot])),
prev_prefix_len: 0,
swa_evicted_seqlen: 0,
swa_branching_seqlen: None,
chunked: false,
priority: 0,
track_adopted_ranges: false,
@@ -591,6 +592,26 @@ fn reinsert_full_backed_target_schedules_mamba_only_backup() {
);
}
#[test]
fn write_back_reinsert_defers_the_mamba_backup_to_eviction() {
let mut tc = mamba_core(/* page_size = */ 1);
tc.set_hicache_enabled();
tc.is_write_back = true;
let key = vec![1, 2];
tc.insert(&insert_params_mamba(&key, &[10, 11], Some(7)));
let leaf = tc.match_prefix(&match_params(&key)).best_match_node_id;
tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new())
.expect("live test node");
let result = tc.insert(&insert_params_mamba(&key, &[20, 21], Some(8)));
assert!(
!result
.cache_actions
.iter()
.any(|action| matches!(action, CacheAction::BackupKV(_)))
);
}
#[test]
fn tombstone_refill_moves_the_node_from_host_to_device_lru() {
let mut tc = mamba_core(/* page_size = */ 1);
rust/sglang-radix-tree/src/tests/components/swa.rs
+390 -17
View File
@@ -549,6 +549,7 @@ fn insert_params_swa<'k>(
mamba_value: None,
prev_prefix_len,
swa_evicted_seqlen,
swa_branching_seqlen: None,
chunked: false,
priority: 0,
track_adopted_ranges: false,
@@ -3669,6 +3670,7 @@ fn build_transfers_are_gated_off_until_the_swa_host_pool_is_wired() {
#[test]
fn backup_host_build_wraps_the_device_value_as_int64() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
tc.set_has_swa_host_pool();
let [a] = chain::<1>(&mut tc);
tc.arena
.set_device_value(a, SWA, Tensor::from_slice(&[5i32]));
@@ -3692,12 +3694,16 @@ fn backup_host_build_wraps_the_device_value_as_int64() {
assert_eq!(device_indices.kind(), Kind::Int64);
assert!(device_indices.equal(&Tensor::from_slice(&[5i64])));
assert!(xfer.host_indices.is_none());
assert!(xfer.nodes_to_load.is_none());
assert_eq!(
xfer.nodes_to_load.as_deref(),
Some(&[tc.arena.node(a).id][..])
);
}
#[test]
fn backup_host_build_returns_none_for_a_tombstone() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
tc.set_has_swa_host_pool();
let [a] = chain::<1>(&mut tc);
let transfers = swa_component(4)
.build_hicache_transfers(
@@ -3717,6 +3723,7 @@ fn backup_host_build_returns_none_for_a_tombstone() {
#[test]
fn backup_spec_reads_the_swa_value_recovered_by_an_earlier_action() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
tc.set_has_swa_host_pool();
let [a] = chain::<1>(&mut tc);
tc.arena
.set_device_value(a, FULL, Tensor::from_slice(&[9i64]));
@@ -3994,26 +4001,47 @@ fn load_back_commit_asserts_the_loaded_length_matches_the_host_indices() {
}
#[test]
fn backup_host_commit_sets_the_host_value_once() {
fn backup_host_commit_sets_the_host_value() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
let [a] = chain::<1>(&mut tc);
set_swa_device(&mut tc, a);
commit_backup(&mut tc, a, &[30i64], /* nodes_to_load = */ None);
assert!(
tc.arena
.host_value(a, SWA)
.equal(&Tensor::from_slice(&[30i64]))
);
}
#[test]
#[should_panic(expected = "is not device-only")]
fn backup_host_commit_rejects_a_target_that_is_already_backed_up() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
let [a] = chain::<1>(&mut tc);
set_swa_device(&mut tc, a);
set_swa_host(&mut tc, a);
let a_id = tc.arena.node(a).id;
commit_backup(&mut tc, a, &[30i64], Some(vec![a_id]));
}
#[test]
#[should_panic(expected = "is not device-only")]
fn backup_host_commit_rejects_a_target_without_a_device_value() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
let [a] = chain::<1>(&mut tc);
let a_id = tc.arena.node(a).id;
commit_backup(&mut tc, a, &[30i64], Some(vec![a_id]));
}
#[test]
fn backup_host_commit_without_offsets_attaches_the_whole_span_once() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
let [a] = chain::<1>(&mut tc);
// A hand-built transfer carries no offsets to scatter by, so the whole
// span attaches to this node and a repeat commit is a no-op.
for host in [30i64, 31] {
let transfer = PoolTransfer {
name: PoolName::Swa,
host_indices: Some(Tensor::from_slice(&[host])),
..Default::default()
};
swa_component(4).commit_hicache_transfer(
&mut tc,
a,
CacheTransferPhase::BackupHost,
vec![transfer],
&mut Vec::new(),
/* insert_result = */ None,
/* pool_storage_result = */ None,
);
commit_backup(&mut tc, a, &[host], /* nodes_to_load = */ None);
}
// The second commit is a no-op: the first host value sticks.
assert!(
tc.arena
.host_value(a, SWA)
@@ -5133,3 +5161,348 @@ fn receipt_anchor_follows_the_locked_node_through_a_split() {
fn new_panics_on_a_zero_sliding_window_size() {
SwaComponent::new(&swa_params_with_window(0));
}
// ==== SWA branching-point caching ====
// A [Full, Swa] core with both the host tier and the SWA host pool wired.
fn swa_hicache_core(window: usize, page_size: usize) -> UnifiedTreeCore<Vec<i64>> {
let mut tc = swa_core(window, page_size);
tc.set_hicache_enabled();
tc.set_has_swa_host_pool();
tc
}
fn set_swa_device_value(tc: &mut UnifiedTreeCore<Vec<i64>>, node: NodeIdx_, value: i64) {
tc.arena
.set_device_value(node, SWA, Tensor::from_slice(&[value]));
}
fn backup_transfers(
tc: &UnifiedTreeCore<Vec<i64>>,
window: usize,
node: NodeIdx_,
) -> Option<Vec<PoolTransfer>> {
swa_component(window)
.build_hicache_transfers(
tc,
node,
CacheTransferPhase::BackupHost,
/* mamba_pool_idx = */ None,
/* host_indices = */ None,
/* token_ids = */ None,
/* prefetch_tokens = */ 0,
/* last_hash = */ None,
)
.unwrap()
}
// The node ids one backup transfer covers, and the device indices it moves.
fn backup_plan(
tc: &UnifiedTreeCore<Vec<i64>>,
window: usize,
node: NodeIdx_,
) -> (Vec<NodeId>, Vec<i64>) {
let transfers = backup_transfers(tc, window, node).expect("a backup transfer");
assert_eq!(transfers.len(), 1);
let xfer = &transfers[0];
assert_eq!(xfer.name, PoolName::Swa);
let device_indices = xfer.device_indices.as_ref().expect("device indices");
assert_eq!(device_indices.kind(), Kind::Int64);
(
xfer.nodes_to_load.clone().expect("nodes_to_load"),
Vec::<i64>::try_from(device_indices).unwrap(),
)
}
fn commit_backup(
tc: &mut UnifiedTreeCore<Vec<i64>>,
node: NodeIdx_,
host_indices: &[i64],
nodes_to_load: Option<Vec<NodeId>>,
) {
swa_component(4).commit_hicache_transfer(
tc,
node,
CacheTransferPhase::BackupHost,
vec![PoolTransfer {
name: PoolName::Swa,
host_indices: Some(Tensor::from_slice(host_indices)),
nodes_to_load,
..Default::default()
}],
&mut Vec::new(),
/* insert_result = */ None,
/* pool_storage_result = */ None,
);
}
#[test]
fn backup_host_build_covers_every_device_only_node_in_the_window() {
let mut tc = swa_hicache_core(/* window = */ 4, /* page_size = */ 1);
let [a, b, c] = chain::<3>(&mut tc);
for (node, value) in [(a, 10i64), (b, 11), (c, 12)] {
set_swa_device_value(&mut tc, node, value);
}
let (nodes, device_indices) = backup_plan(&tc, /* window = */ 4, c);
// Ancestors first, so the publish side links each store event to its parent.
let expected = [a, b, c].map(|node| tc.arena.node(node).id);
assert_eq!(nodes, expected.to_vec());
assert_eq!(device_indices, vec![10, 11, 12]);
}
#[test]
fn backup_host_build_stops_at_a_node_another_ack_owns() {
let mut tc = swa_hicache_core(/* window = */ 4, /* page_size = */ 1);
let [a, b, c] = chain::<3>(&mut tc);
for (node, value) in [(a, 10i64), (b, 11), (c, 12)] {
set_swa_device_value(&mut tc, node, value);
}
let b_id = tc.arena.node(b).id;
tc.mark_write_through_pending(vec![b_id], /* ack_id = */ b_id)
.expect("live test node");
// `b`'s ack already owns `b` and everything above it, so this backup takes
// only what is left below it: two acks never claim the same node.
let (nodes, device_indices) = backup_plan(&tc, /* window = */ 4, c);
assert_eq!(nodes, vec![tc.arena.node(c).id]);
assert_eq!(device_indices, vec![12]);
}
#[test]
fn backup_host_build_stops_at_the_sliding_window_edge() {
let mut tc = swa_hicache_core(/* window = */ 2, /* page_size = */ 1);
let [a, b, c] = chain::<3>(&mut tc);
for (node, value) in [(a, 10i64), (b, 11), (c, 12)] {
set_swa_device_value(&mut tc, node, value);
}
let (nodes, device_indices) = backup_plan(&tc, /* window = */ 2, c);
assert_eq!(nodes, vec![tc.arena.node(b).id, tc.arena.node(c).id]);
assert_eq!(device_indices, vec![11, 12]);
}
#[test]
fn backup_host_build_walks_past_a_node_that_is_already_backed_up() {
let mut tc = swa_hicache_core(/* window = */ 4, /* page_size = */ 1);
let [a, b, c] = chain::<3>(&mut tc);
for (node, value) in [(a, 10i64), (b, 11), (c, 12)] {
set_swa_device_value(&mut tc, node, value);
}
set_swa_host(&mut tc, b);
// `b` is backed up already: it consumes window budget but is not re-sent,
// and the walk continues to the unbacked ancestor above it.
let (nodes, device_indices) = backup_plan(&tc, /* window = */ 4, c);
assert_eq!(nodes, vec![tc.arena.node(a).id, tc.arena.node(c).id]);
assert_eq!(device_indices, vec![10, 12]);
}
#[test]
fn backup_host_build_is_none_without_an_swa_host_pool() {
let mut tc = swa_core(/* window = */ 4, /* page_size = */ 1);
let [a] = chain::<1>(&mut tc);
set_swa_device_value(&mut tc, a, 10);
assert!(backup_transfers(&tc, /* window = */ 4, a).is_none());
}
#[test]
fn backup_host_commit_scatters_the_host_span_across_the_covered_nodes() {
let mut tc = swa_hicache_core(/* window = */ 4, /* page_size = */ 1);
// Spans 1 / 2 / 1: the scatter has to advance by each node's own length.
let mut parent = tc.arena.root();
let mut nodes = Vec::new();
for key in [vec![1i64], vec![2, 3], vec![4]] {
parent = tc
.arena
.alloc_child(
parent, key, /* priority = */ 0, /* extra_key = */ None,
)
.unwrap();
nodes.push(parent);
}
let [a, b, c] = [nodes[0], nodes[1], nodes[2]];
for (node, value) in [(a, vec![10i64]), (b, vec![11, 12]), (c, vec![13])] {
tc.arena
.set_device_value(node, SWA, Tensor::from_slice(&value));
}
let (ids, device_indices) = backup_plan(&tc, /* window = */ 4, c);
assert_eq!(device_indices, vec![10, 11, 12, 13]);
commit_backup(&mut tc, c, &[100i64, 101, 102, 103], Some(ids));
for (node, host) in [(a, vec![100i64]), (b, vec![101, 102]), (c, vec![103])] {
assert!(
tc.arena
.host_value(node, SWA)
.equal(&Tensor::from_slice(&host))
);
}
}
#[test]
fn needs_incremental_backup_tracks_the_unbacked_window() {
let mut tc = swa_hicache_core(/* window = */ 4, /* page_size = */ 1);
let [a, b] = chain::<2>(&mut tc);
let swa = swa_component(4);
assert!(!TreeComponent::<Vec<i64>>::needs_incremental_backup(
&swa, &tc, b
));
// A device-only ancestor is enough, even when the target itself is clean.
set_swa_device_value(&mut tc, a, 10);
set_swa_device_value(&mut tc, b, 11);
set_swa_host(&mut tc, b);
assert!(TreeComponent::<Vec<i64>>::needs_incremental_backup(
&swa, &tc, b
));
set_swa_host(&mut tc, a);
assert!(!TreeComponent::<Vec<i64>>::needs_incremental_backup(
&swa, &tc, b
));
}
#[test]
fn buffer_mode_backup_window_is_the_target_alone_at_both_call_sites() {
// Cache mode walks the window: a host-only target still reaches the
// device-only ancestor above it. Buffer mode stages the target alone.
let mut tc = swa_hicache_core(/* window = */ 4, /* page_size = */ 1);
let [a, b] = chain::<2>(&mut tc);
set_swa_device_value(&mut tc, a, 10);
set_swa_host(&mut tc, b);
let swa = swa_component(4);
let a_id = tc.arena.node(a).id;
let b_id = tc.arena.node(b).id;
assert!(TreeComponent::<Vec<i64>>::needs_incremental_backup(
&swa, &tc, b
));
assert_eq!(
backup_transfers(&tc, 4, b).unwrap()[0].nodes_to_load,
Some(vec![a_id])
);
tc.set_host_memory_buffer_only();
assert!(!TreeComponent::<Vec<i64>>::needs_incremental_backup(
&swa, &tc, b
));
assert!(backup_transfers(&tc, 4, b).is_none());
// A device-resident target is staged by itself, whatever its ancestors hold.
set_swa_host(&mut tc, a);
set_swa_device_value(&mut tc, b, 11);
assert!(TreeComponent::<Vec<i64>>::needs_incremental_backup(
&swa, &tc, b
));
assert_eq!(
backup_transfers(&tc, 4, b).unwrap()[0].nodes_to_load,
Some(vec![b_id])
);
}
#[test]
fn write_back_reinsert_still_backs_up_an_unbacked_swa_window() {
let mut tc: UnifiedTreeCore<Vec<i64>> = UnifiedTreeCore::new(
CacheInitParams {
is_write_back: true,
enable_hicache: true,
has_swa_host_pool: true,
..swa_params_with_window(4)
},
vec![FULL, SWA],
);
let key = vec![1, 2];
tc.insert(&insert_params_swa(&key, &[10, 11], 0, 0));
let leaf_idx = child_of(&tc, tc.arena.root(), &[1]);
let leaf = tc.arena.node(leaf_idx).id;
// The cache applied the SwaRebuild; Full is on host, SWA is still device-only.
store_swa_device(&mut tc, leaf_idx);
tc.commit_backup(leaf, Tensor::from_slice(&[100i64, 101]), HashMap::new())
.expect("live test node");
let result = tc.insert(&insert_params_swa(&key, &[20, 21], 0, 0));
let backups: Vec<_> = result
.cache_actions
.iter()
.filter_map(|action| match action {
CacheAction::BackupKV(backup) => Some(backup.node_ids.clone()),
_ => None,
})
.collect();
assert_eq!(backups, vec![vec![leaf]]);
let (full_device_indices, comp_xfers) = tc.build_backup_spec(leaf).expect("live test node");
assert_eq!(full_device_indices.numel(), 0);
assert_eq!(comp_xfers[&SWA][0].nodes_to_load, Some(vec![leaf]));
}
// Finalize an otherwise-empty match carrying the given Full-KV reach.
fn finalize_branching(
tc: &UnifiedTreeCore<Vec<i64>>,
device_len: usize,
host_hit_length: usize,
full_kv_hit_length: usize,
) -> Option<usize> {
swa_component(4)
.finalize_match_result_in_tree_core(
tc,
MatchResult {
device_indices: Tensor::from_slice(&vec![0i64; device_len]),
host_hit_length,
full_kv_hit_length,
..tc.empty_match_result()
},
/* last_device_node_idx = */ tc.arena.root(),
/* best_match_node_idx = */ tc.arena.root(),
&MatchPrefixParams {
key: &Vec::new(),
namespace: Default::default(),
},
&[],
0,
)
.swa_branching_seqlen
}
#[test]
fn match_reports_the_page_aligned_swa_branching_seqlen() {
let tc = swa_hicache_core(/* window = */ 4, /* page_size = */ 4);
// Full KV reaches 11 tokens, the SWA boundary only 2; 11 aligns down to 8.
assert_eq!(
finalize_branching(&tc, /* device = */ 2, /* host_hit = */ 0, 11),
Some(8)
);
// Host-loaded Full KV counts toward the boundary the branch must beat.
assert_eq!(
finalize_branching(&tc, /* device = */ 2, /* host_hit = */ 6, 11),
None
);
}
#[test]
fn swa_branching_seqlen_is_none_when_no_aligned_page_lies_beyond_the_window() {
let tc = swa_hicache_core(/* window = */ 4, /* page_size = */ 4);
// Full KV reaches 3 tokens, which aligns down to 0: nothing to branch at.
assert_eq!(
finalize_branching(&tc, /* device = */ 0, /* host_hit = */ 0, 3),
None
);
// The aligned position must lie strictly beyond the boundary.
assert_eq!(
finalize_branching(&tc, /* device = */ 8, /* host_hit = */ 0, 11),
None
);
}
#[test]
fn insert_reports_whether_it_reached_the_branch_boundary() {
for (branching_seqlen, expected) in [(Some(3), true), (Some(4), false), (None, false)] {
let mut tc = swa_hicache_core(/* window = */ 4, /* page_size = */ 1);
let result = tc.insert(&InsertParams {
swa_branching_seqlen: branching_seqlen,
..insert_params_swa(&vec![1, 2, 3], &[10, 11, 12], 0, 0)
});
assert_eq!(
result.swa_branch_inserted, expected,
"swa_branching_seqlen={branching_seqlen:?}"
);
}
}
rust/sglang-radix-tree/src/tests/unified_tree_core.rs
+3
View File
@@ -1863,6 +1863,7 @@ fn insert_params<'k>(key: &'k Vec<i64>, value: &[i64]) -> InsertParams<'k, Vec<i
mamba_value: None,
prev_prefix_len: 0,
swa_evicted_seqlen: 0,
swa_branching_seqlen: None,
chunked: false,
priority: 0,
track_adopted_ranges: false,
@@ -3051,6 +3052,7 @@ fn bigram_insert_events_carry_pair_token_payloads() {
mamba_value: None,
prev_prefix_len: 0,
swa_evicted_seqlen: 0,
swa_branching_seqlen: None,
chunked: false,
priority: 0,
track_adopted_ranges: false,
@@ -8398,6 +8400,7 @@ fn sequence_insert_params<'k>(
mamba_value,
prev_prefix_len,
swa_evicted_seqlen: 0,
swa_branching_seqlen: None,
chunked: false,
priority: 0,
track_adopted_ranges: false,
rust/sglang-radix-tree/src/unified_tree_core.rs
+29 -2
View File
@@ -92,6 +92,8 @@ pub struct MatchResult {
/// SWA tokens that hit on host (within the sliding window) and will be
/// loaded back into the SWA device pool.
pub swa_host_hit_length: usize,
/// The longest page-aligned position that could have hit if an SWA window existed.
pub swa_branching_seqlen: Option<usize>,
/// Mamba slots that hit on host and will be loaded back; 0 or 1.
pub mamba_host_hit_length: usize,
/// The longest chunk-aligned position that could have hit if a mamba state existed.
@@ -124,6 +126,8 @@ pub struct InsertParams<'k, K: ChildKeyType> {
pub prev_prefix_len: usize,
/// The request's SWA-evicted prefix boundary; SWA data below it stays tombstoned.
pub swa_evicted_seqlen: usize,
/// The Full-KV-derived boundary whose SWA window this insert should materialize.
pub swa_branching_seqlen: Option<usize>,
/// The donated mamba slot for the insert target leaf; None on non-mamba trees.
pub mamba_value: Option<Tensor>,
/// Whether this is a chunked-prefill insert (no hit-count bump).
@@ -146,6 +150,8 @@ pub struct InsertResult {
/// Whether the cache holds Mamba state covering the inserted sequence;
/// vacuously true for an empty insert.
pub mamba_exist: bool,
/// Whether this insert reached the requested SWA branch boundary.
pub swa_branch_inserted: bool,
/// The deepest host-backed node an insert_host attached or matched.
pub inserted_host_node: Option<NodeId>,
/// Whether write-through rejected a host suffix below an unbacked parent.
@@ -207,6 +213,7 @@ pub struct InsertWalkState<K: ChildKeyType> {
namespace: KeyNamespace,
prev_prefix_len: usize,
swa_evicted_seqlen: usize,
swa_branching_seqlen: Option<usize>,
mamba_value: Option<Tensor>,
chunked: bool,
priority: i64,
@@ -532,6 +539,9 @@ pub struct UnifiedTreeCore<K: ChildKeyType> {
pub(crate) is_write_back: bool,
/// Whether the host tier (HiCache) is wired.
pub(crate) enable_hicache: bool,
/// Whether the host tier stages one node per FIFO backup intent rather than
/// caching windows; buffer mode sizes its storage keys off a single node.
pub(crate) is_host_memory_buffer_only: bool,
/// Whether the storage tier (L3) is wired; gates page-hash computation.
pub(crate) enable_storage: bool,
/// Whether a direct device-to-external-cache linker is wired.
@@ -717,6 +727,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
page_size: params.page_size,
is_write_back: params.is_write_back,
enable_hicache: params.enable_hicache,
is_host_memory_buffer_only: false,
enable_storage: false,
enable_external_cache_linker: false,
has_swa_host_pool: params.has_swa_host_pool,
@@ -1228,6 +1239,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
mamba_host_hit_length: 0,
mamba_branching_seqlen: None,
swa_host_hit_length: 0,
swa_branching_seqlen: None,
full_kv_hit_length,
cache_actions: Vec::new(),
};
@@ -1257,6 +1269,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
best_match_node_id: root_id,
host_hit_length: 0,
swa_host_hit_length: 0,
swa_branching_seqlen: None,
full_kv_hit_length: 0,
mamba_host_hit_length: 0,
mamba_branching_seqlen: None,
@@ -1389,6 +1402,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
inserted_host_node: None,
host_insert_dropped: false,
mamba_exist: true,
swa_branch_inserted: false,
adopted_ranges: None,
cache_actions: Vec::new(),
}),
@@ -1410,6 +1424,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
namespace: params.namespace.to_owned(),
prev_prefix_len: params.prev_prefix_len,
swa_evicted_seqlen: params.swa_evicted_seqlen,
swa_branching_seqlen: params.swa_branching_seqlen,
mamba_value: params.mamba_value.as_ref().map(Tensor::shallow_clone),
chunked: params.chunked,
priority: params.priority,
@@ -1535,6 +1550,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
value: state.value.shallow_clone(),
prev_prefix_len: state.prev_prefix_len,
swa_evicted_seqlen: state.swa_evicted_seqlen,
swa_branching_seqlen: state.swa_branching_seqlen,
mamba_value: state.mamba_value.as_ref().map(Tensor::shallow_clone),
chunked: state.chunked,
priority: state.priority,
@@ -1685,6 +1701,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
value: state.value.shallow_clone(),
prev_prefix_len: state.prev_prefix_len,
swa_evicted_seqlen: state.swa_evicted_seqlen,
swa_branching_seqlen: state.swa_branching_seqlen,
mamba_value: state.mamba_value.as_ref().map(Tensor::shallow_clone),
chunked: state.chunked,
priority: state.priority,
@@ -1707,7 +1724,11 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
/// Whether an auxiliary component has new device data missing from Host.
fn needs_incremental_component_backup_(&self, node_id: NodeIdx_) -> bool {
self.components.iter().any(|component| {
component.component_type() != BASE_COMPONENT_TYPE
let component_type = component.component_type();
// Write-back defers Full and Mamba to eviction; SWA still publishes
// its window here because out-of-window frees never reach eviction.
component_type != BASE_COMPONENT_TYPE
&& (!self.is_write_back || component_type == SWA)
&& component.needs_incremental_backup(self, node_id)
})
}
@@ -1724,7 +1745,6 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
let node = self.arena.node(target_node_id);
self.enable_hicache
&& !self.is_write_back
&& node.backuped()
&& node.write_through_pending_id.is_none()
&& self.needs_incremental_component_backup_(target_node_id)
@@ -2700,6 +2720,11 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
self.enable_hicache = true;
}
/// Mark the host tier as buffer-only; wired after the host pools are built.
pub fn set_host_memory_buffer_only(&mut self) {
self.is_host_memory_buffer_only = true;
}
/// Whether the storage tier (L3) is wired; storage attaches after tree construction.
pub fn set_enable_storage(&mut self, value: bool) {
self.enable_storage = value;
@@ -2983,6 +3008,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
inserted_host_node: None,
host_insert_dropped: false,
mamba_exist: true,
swa_branch_inserted: false,
adopted_ranges: None,
cache_actions: Vec::new(),
});
@@ -3022,6 +3048,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
inserted_host_node: None,
host_insert_dropped: false,
mamba_exist: false,
swa_branch_inserted: false,
adopted_ranges: None,
cache_actions,
};