minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 11 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
python/sglang/srt/mem_cache/rust_tree_core/adapter.py
+10
View File
@@ -228,6 +228,7 @@ def _insert_step_from_binding(step) -> InsertStepResult:
prefix_len=step.result.prefix_len,
last_device_node=step.result.last_device_node,
mamba_exist=step.result.mamba_exist,
swa_branch_inserted=step.result.swa_branch_inserted,
host_insert_dropped=step.result.host_insert_dropped,
adopted_ranges=(
{
@@ -252,6 +253,7 @@ def _match_result_from_binding(result) -> MatchResult:
best_match_node=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,
@@ -589,6 +591,7 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
mamba_value=params.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=0 if params.priority is None else params.priority,
track_adopted_ranges=params.track_adopted_ranges,
@@ -630,6 +633,13 @@ class RustUnifiedTreeCore(UnifiedTreeCoreInterface):
def set_hicache_enabled(self) -> None:
self._binding.set_hicache_enabled()
def set_host_memory_buffer_only(self) -> None:
self._binding.set_host_memory_buffer_only()
@property
def is_host_memory_buffer_only(self) -> bool:
return self._binding.is_host_memory_buffer_only()
@property
def page_size(self) -> int:
# Read-only: the Rust core freezes it at construction.
python/sglang/srt/mem_cache/unified_cache/components/swa_component.py
+30 -17
View File
@@ -89,13 +89,27 @@ class SWAComponent(TreeComponent):
component_type = ComponentType.SWA
def _dirty_backup_window(self, node: UnifiedTreeNode) -> list[UnifiedTreeNode]:
def _collect_unbacked_swa_nodes(
self, node: UnifiedTreeNode
) -> list[UnifiedTreeNode]:
"""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``.
"""
if not self.tree_core.has_swa_host_pool:
return []
if self.tree_core.is_host_memory_buffer_only:
cd = node.component_data[self.component_type]
return [node] if cd.value is not None else []
return self._collect_unbacked_swa_nodes_in_window(node)
def _collect_unbacked_swa_nodes_in_window(
self, node: UnifiedTreeNode
) -> list[UnifiedTreeNode]:
ct = self.component_type
covered = 0
dirty: list[UnifiedTreeNode] = []
unbacked: list[UnifiedTreeNode] = []
cur = node
while (
cur is not self.tree_core.root_node and covered < self.sliding_window_size
@@ -109,12 +123,12 @@ class SWAComponent(TreeComponent):
break
covered += len(value)
if cd.value is not None and cd.host_value is None:
dirty.append(cur)
unbacked.append(cur)
cur = cur.parent
return dirty
return unbacked
def needs_incremental_backup(self, node: UnifiedTreeNode) -> bool:
return bool(self._dirty_backup_window(node))
return bool(self._collect_unbacked_swa_nodes(node))
def reset_session_state(self) -> None:
super().reset_session_state()
@@ -919,6 +933,11 @@ class SWAComponent(TreeComponent):
# that boundary so insertion creates a tombstone instead of live SWA KV.
insert_params.swa_evicted_seqlen = req.kv.swa_evicted_seqlen
# A recurrent checkpoint must stay attached to its exact token prefix.
# Let MambaComponent select the insertion length for hybrid caches.
if self.cache.is_mamba_enabled:
return None
branching_seqlen = req.swa_branching_seqlen
if branching_seqlen is None or branching_seqlen <= req.kv.cache_protected_len:
return None
@@ -991,8 +1010,7 @@ class SWAComponent(TreeComponent):
elif prefetch_pages <= 0:
return PreparePrefetchResult()
elif (
self.tree_core.is_root(node_id)
or self.cache.host_memory_mode == "buffer_only"
self.tree_core.is_root(node_id) or self.tree_core.is_host_memory_buffer_only
):
# Sub-window fetch: at root the sequence IS its window; mid-tree
# (buffer mode) the window head is the device prefix's own ring
@@ -1030,22 +1048,17 @@ class SWAComponent(TreeComponent):
return None
if phase == CacheTransferPhase.BACKUP_HOST:
if self.cache.host_memory_mode == "buffer_only":
# Buffer mode stages one node/hash span per FIFO backup intent.
cd = node.component_data[ct]
dirty = [node] if cd.value is not None else []
else:
dirty = self._dirty_backup_window(node)
if not dirty:
unbacked_swa_nodes = self._collect_unbacked_swa_nodes(node)
if not unbacked_swa_nodes:
return None
dirty.reverse()
unbacked_swa_nodes.reverse()
return [
PoolTransfer(
name=PoolName.SWA,
device_indices=torch.cat(
[n.component_data[ct].value for n in dirty]
[n.component_data[ct].value for n in unbacked_swa_nodes]
).to(torch.int64),
nodes_to_load=[n.id for n in dirty],
nodes_to_load=[n.id for n in unbacked_swa_nodes],
)
]
python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py
+4
View File
@@ -415,6 +415,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
self.page_size = params.page_size
self.is_eagle = params.is_eagle and ComponentType.MAMBA not in components
self.enable_hicache = False
self.is_host_memory_buffer_only = False
self.enable_storage = False
self.enable_external_cache_linker = False
self.write_through_threshold = 256
@@ -2045,6 +2046,9 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
def set_hicache_enabled(self) -> None:
self.enable_hicache = True
def set_host_memory_buffer_only(self) -> None:
self.is_host_memory_buffer_only = True
def insert_host(
self,
node_id: NodeId,
python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py
+7
View File
@@ -152,6 +152,8 @@ class UnifiedTreeCoreInterface(ABC):
write_through_threshold: int
is_write_back: bool
has_swa_host_pool: bool
# Whether the host tier stages one node per FIFO backup intent.
is_host_memory_buffer_only: bool
kv_events: KVCacheEventRecorder
# ==== Tree API ====
@@ -456,6 +458,11 @@ class UnifiedTreeCoreInterface(ABC):
"""Mark the host tier (HiCache) as wired."""
...
@abstractmethod
def set_host_memory_buffer_only(self) -> None:
"""Mark the host tier as buffer-only: one node staged per backup intent."""
...
@abstractmethod
def insert_host(
self,
python/sglang/srt/mem_cache/unified_radix_cache.py
+1
View File
@@ -456,6 +456,7 @@ class UnifiedRadixCache(BasePrefixCache):
self.tree_core.has_swa_host_pool = swa._swa_kv_pool_host is not None
if self.host_memory_mode == "buffer_only":
self.tree_core.set_host_memory_buffer_only()
swa = self.components.get(ComponentType.SWA)
validate_buffer_only_stack(
sidecar_pool_specs=self.sidecar_pool_specs,
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,
};
test/registered/unit/mem_cache/test_rust_tree_core_integration.py
+25
View File
@@ -101,6 +101,7 @@ def _pump_insert(core: RustUnifiedTreeCore, params: InsertParams) -> InsertResul
prefix_len=step.result.prefix_len,
last_device_node=step.result.last_device_node,
mamba_exist=step.result.mamba_exist,
swa_branch_inserted=step.result.swa_branch_inserted,
cache_actions=actions,
)
@@ -2288,5 +2289,29 @@ def test_stale_inspection_handles_raise_key_error_or_report_absence():
assert not core.is_node_in_host_lru(stale_root, ComponentType.SWA)
# ---- SWA branching-point caching ----
def _swa_hicache_core(window: int = 8) -> RustUnifiedTreeCore:
core = _swa_tree_core(window=window)
core.set_hicache_enabled()
core.has_swa_host_pool = True
return core
def test_insert_reports_whether_it_reached_the_swa_branch_boundary():
for branching_seqlen, expected in [(2, True), (3, False), (None, False)]:
core = _swa_hicache_core()
result = _pump_insert(
core,
InsertParams(
key=_key([1, 2]),
value=torch.tensor([10, 11], dtype=torch.int64),
swa_branching_seqlen=branching_seqlen,
),
)
assert result.swa_branch_inserted is expected, branching_seqlen
if __name__ == "__main__":
sys.exit(pytest.main([__file__]))
test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py
+78 -58
View File
@@ -5628,6 +5628,44 @@ class UnifiedRadixCacheSuite:
cache.tree_core.has_swa_host_pool, swa._swa_kv_pool_host is not None
)
def test_swa_backup_collector_is_shared_by_both_call_sites(self):
"""needs_incremental_backup and the BACKUP_HOST transfer read one
collector: cache mode walks the window past a host-backed target to
its device-only ancestor, buffer mode stages the target alone."""
if not self.cfg.has_swa or self.cfg.has_mamba:
self.skipTest("requires SWA-only")
if self.cfg.sliding_window_size <= self.cfg.page_size:
self.skipTest("the window must reach past the leaf's own page")
if _selected_tree_core_test_backend() == "rust":
# needs_incremental_backup is a component method on Python nodes;
# the Rust core pins the same contract in its own unit suite.
self.skipTest("component-level check is Python-core only")
cache, allocator, req_to_token_pool = self._build_hicache_fixture()
chain = self._build_chain_pages(cache, allocator, req_to_token_pool, 2)
if len(chain) < 2:
self.skipTest("chain too short")
parent, leaf = chain[-2], chain[-1]
swa = cache.components[ComponentType.SWA]
leaf_node = cache.tree_core.node_by_id(leaf)
cache.tree_core.set_component_host_value_raw(
leaf,
ComponentType.SWA,
_device_value(cache, leaf, ComponentType.SWA).clone(),
)
self.assertTrue(swa.needs_incremental_backup(leaf_node))
xfer = cache.tree_core.build_hicache_transfers(
ComponentType.SWA, leaf, CacheTransferPhase.BACKUP_HOST
)[0]
self.assertEqual(xfer.nodes_to_load, [parent])
cache.tree_core.set_host_memory_buffer_only()
self.assertTrue(swa.needs_incremental_backup(leaf_node))
xfer = cache.tree_core.build_hicache_transfers(
ComponentType.SWA, leaf, CacheTransferPhase.BACKUP_HOST
)[0]
self.assertEqual(xfer.nodes_to_load, [leaf])
def test_zero_match_result_carries_node_id_handles(self):
cache, allocator, req_to_token_pool = build_fixture(self.cfg)
ps = self.cfg.page_size
@@ -5872,14 +5910,7 @@ class UnifiedRadixCacheSuite:
self.assertEqual(result.host_hit_length, 0)
self.assertEqual(result.swa_host_hit_length, _node_key_length(cache, leaf))
def _skip_swa_branching_on_rust(self) -> None:
# TODO(alphabetc1): drop this gate once #37584 ports SWA branching-point
# caching to the Rust tree core.
if _selected_tree_core_test_backend() == "rust":
self.skipTest("SWA branching-point caching is Python-core only")
def test_swa_branching_seqlen_uses_device_full_hit(self):
self._skip_swa_branching_on_rust()
if (
not self.cfg.has_swa
or self.cfg.has_mamba
@@ -5895,20 +5926,13 @@ class UnifiedRadixCacheSuite:
self._insert(cache, allocator, req_to_token_pool, prefix)
self._insert(cache, allocator, req_to_token_pool, tokens)
leaf = cache.resolve_node_handle(
cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)))
).last_device_node
leaf = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)))
).last_device_node
evicted = cache.tree_core.evict_component(
leaf, ComponentType.SWA, EvictLayer.DEVICE
)
device_frees = defaultdict(list)
cache.tree_core._evict_component_and_detach_lru(
leaf,
cache.components[ComponentType.SWA],
device_frees=device_frees,
host_frees=defaultdict(list),
target=EvictLayer.DEVICE,
)
cache._drain_device_frees(device_frees)
cache._free_values(evicted.device_frees, evicted.host_frees)
result = cache.match_prefix(MatchPrefixParams(key=RadixKey(array("q", tokens))))
@@ -5931,7 +5955,6 @@ class UnifiedRadixCacheSuite:
self.assertIsNone(rematch.swa_branching_seqlen)
def test_swa_branching_seqlen_uses_host_full_hit(self):
self._skip_swa_branching_on_rust()
if (
not self.cfg.has_swa
or self.cfg.has_mamba
@@ -5946,24 +5969,21 @@ class UnifiedRadixCacheSuite:
self._insert(cache, allocator, req_to_token_pool, prefix)
self._insert(cache, allocator, req_to_token_pool, tokens)
leaf = cache.resolve_node_handle(
cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)))
).last_device_node
)
parent = leaf.parent
self._backup_node(cache, leaf.id)
lock_result = cache.inc_lock_ref(parent.id)
leaf = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", tokens)))
).last_device_node
leaf_len = _node_key_length(cache, leaf)
parent = _node_parent(cache, leaf)
self._backup_node(cache, leaf)
lock_result = cache.inc_lock_ref(parent)
try:
cache.evict(EvictParams(num_tokens=len(leaf.key)))
cache.evict(EvictParams(num_tokens=leaf_len))
finally:
cache.dec_lock_ref(parent.id, lock_result.to_dec_params())
device_frees = defaultdict(list)
host_frees = defaultdict(list)
cache.components[ComponentType.SWA].evict_component(
leaf, device_frees, host_frees, target=EvictLayer.HOST
cache.dec_lock_ref(parent, lock_result.to_dec_params())
evicted = cache.tree_core.evict_component(
leaf, ComponentType.SWA, EvictLayer.HOST
)
cache._free_values(device_frees, host_frees)
cache._free_values(evicted.device_frees, evicted.host_frees)
full_host_pool = cache.cache_controller.mem_pool_host
swa_host_pool = cache.components[ComponentType.SWA]._swa_kv_pool_host
full_available_before = full_host_pool.available_size()
@@ -5985,7 +6005,7 @@ class UnifiedRadixCacheSuite:
self.assertEqual(full_host_pool.available_size(), full_available_before)
self.assertEqual(
swa_host_pool.available_size(),
swa_available_before - len(leaf.key),
swa_available_before - leaf_len,
)
rematch = cache.match_prefix(
@@ -6799,7 +6819,6 @@ class UnifiedRadixCacheSuite:
self.assertEqual(comp_xfers[ComponentType.SWA][0].nodes_to_load, [a, b])
def test_hicache_swa_backup_window_stops_at_pending_ancestor(self):
self._skip_swa_branching_on_rust()
if (
not self.cfg.has_swa
or self.cfg.has_mamba
@@ -6816,24 +6835,30 @@ class UnifiedRadixCacheSuite:
c = chain[-1]
c_swa = _device_value(cache, c, ComponentType.SWA).clone()
# First transfer: publish Full for C only, leaving SWA dirty while the
# write-through ack is still pending.
# First transfer: publish Full for C only. C's SWA is a device tombstone
# (decode-evicted, never backed up), so the write-through ack stays
# pending on a node the SWA backup window has nothing to send for.
cache.tree_core.set_component_device_value_raw(c, ComponentType.SWA, None)
if cache.tree_core.is_node_in_device_lru(c, ComponentType.SWA):
cache.tree_core.remove_node_from_device_lru(c, ComponentType.SWA)
cache.tree_core.set_component_evictable_size(
ComponentType.SWA,
cache.tree_core.component_evictable_size(ComponentType.SWA) - len(c_swa),
)
self.assertGreater(
cache._execute_and_commit_kv_backup(BackupKV(node_ids=[c])),
0,
)
self.assertEqual(
cache.tree_core.node_by_id(c).write_through_pending_id,
c,
)
self.assertEqual(cache.tree_core.get_write_through_pending_id(c), c)
self.assertIsNotNone(_host_value(cache, c, ComponentType.FULL))
self.assertIsNone(_host_value(cache, c, ComponentType.SWA))
# Simulate SWA being reconstructed on device before the first ack. The
# next incremental SWA backup must treat C as the boundary and back up
# only the newly inserted descendant.
cache.tree_core.set_component_device_value_raw(c, ComponentType.SWA, c_swa)
# SWA is reconstructed on device before the first ack, the way a
# load-back commit stores it: under the pending segment lock the value
# counts as protected until the ack releases it. The next incremental
# SWA backup must treat C as the boundary and back up only the newly
# inserted descendant.
cache.tree_core.set_component_device_value(c, ComponentType.SWA, c_swa)
tokens = self._match_tokens_for_chain(cache, chain)
next_tokens = tokens + self._make_seq(9000, 1)
cache.write_through_threshold = 1
@@ -6843,20 +6868,15 @@ class UnifiedRadixCacheSuite:
d = cache.match_prefix(
MatchPrefixParams(key=RadixKey(array("q", next_tokens)))
).last_device_node
self.assertEqual(
cache.tree_core.node_by_id(c).write_through_pending_id,
c,
)
self.assertEqual(
cache.tree_core.node_by_id(d).write_through_pending_id,
d,
)
self.assertEqual(cache.tree_core.get_write_through_pending_id(c), c)
self.assertEqual(cache.tree_core.get_write_through_pending_id(d), d)
self.assertIsNone(_host_value(cache, c, ComponentType.SWA))
self.assertIsNotNone(_host_value(cache, d, ComponentType.SWA))
cache.writing_check(write_back=True)
self.assertIsNone(cache.tree_core.node_by_id(c).write_through_pending_id)
self.assertIsNone(cache.tree_core.node_by_id(d).write_through_pending_id)
self.assertIsNone(cache.tree_core.get_write_through_pending_id(c))
self.assertIsNone(cache.tree_core.get_write_through_pending_id(d))
cache.sanity_check()
def _swa_finalize_setup(self):
"""Build a SWA chain long enough to fill at least the window