From bd45cd50ca900dd821f829ca9adfbf9aa3336bda Mon Sep 17 00:00:00 2001 From: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com> Date: Sat, 12 Sep 2026 18:37:57 +0800 Subject: [PATCH] [Unified Tree] Port SWA Branching-Point Caching to the Rust TreeCore (#37584) --- .../srt/mem_cache/rust_tree_core/adapter.py | 10 + .../unified_cache/components/swa_component.py | 47 +- .../unified_cache/unified_tree_core.py | 4 + .../unified_tree_core_interface.py | 7 + .../srt/mem_cache/unified_radix_cache.py | 1 + rust/sglang-radix-tree/src/components/swa.rs | 141 +++++- rust/sglang-radix-tree/src/python_bindings.rs | 35 +- .../src/tests/components/base.rs | 1 + .../src/tests/components/full.rs | 2 + .../src/tests/components/mamba.rs | 21 + .../src/tests/components/swa.rs | 407 +++++++++++++++++- .../src/tests/unified_tree_core.rs | 3 + .../src/unified_tree_core.rs | 31 +- .../test_rust_tree_core_integration.py | 25 ++ .../test_unified_radix_cache_unittest.py | 136 +++--- 15 files changed, 762 insertions(+), 109 deletions(-) diff --git a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py index 071bb9134..ff11a09a8 100644 --- a/python/sglang/srt/mem_cache/rust_tree_core/adapter.py +++ b/python/sglang/srt/mem_cache/rust_tree_core/adapter.py @@ -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. diff --git a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py index 1c2c95156..49f306458 100644 --- a/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py +++ b/python/sglang/srt/mem_cache/unified_cache/components/swa_component.py @@ -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], ) ] diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py index 8e0555f5f..3ee02de12 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py @@ -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, diff --git a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py index 64c36443a..4fabcd2dc 100644 --- a/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py +++ b/python/sglang/srt/mem_cache/unified_cache/unified_tree_core_interface.py @@ -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, diff --git a/python/sglang/srt/mem_cache/unified_radix_cache.py b/python/sglang/srt/mem_cache/unified_radix_cache.py index 55d89437d..2f44b0568 100644 --- a/python/sglang/srt/mem_cache/unified_radix_cache.py +++ b/python/sglang/srt/mem_cache/unified_radix_cache.py @@ -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, diff --git a/rust/sglang-radix-tree/src/components/swa.rs b/rust/sglang-radix-tree/src/components/swa.rs index 84525821a..4789adbe3 100644 --- a/rust/sglang-radix-tree/src/components/swa.rs +++ b/rust/sglang-radix-tree/src/components/swa.rs @@ -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_( + &self, + tree_core: &UnifiedTreeCore, + node_id: NodeIdx_, + ) -> Vec { + 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_( + &self, + tree_core: &UnifiedTreeCore, + node_id: NodeIdx_, + ) -> Vec { + let mut covered_tokens = 0; + let mut unbacked: Vec = 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( tree_core: &UnifiedTreeCore, from: Option, @@ -305,6 +359,12 @@ impl SwaComponent { } impl TreeComponent for SwaComponent { + fn needs_incremental_backup(&self, tree_core: &UnifiedTreeCore, 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 TreeComponent 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 TreeComponent for SwaComponent { result: &mut InsertResult, cache_actions: &mut Vec, ) { + 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 TreeComponent 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, Vec) = + 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 TreeComponent 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 diff --git a/rust/sglang-radix-tree/src/python_bindings.rs b/rust/sglang-radix-tree/src/python_bindings.rs index 9a81f2dc1..d34a2862b 100644 --- a/rust/sglang-radix-tree/src/python_bindings.rs +++ b/rust/sglang-radix-tree/src/python_bindings.rs @@ -301,6 +301,8 @@ struct InspectionMatchResultInput { #[pyo3(attribute)] swa_host_hit_length: usize, #[pyo3(attribute)] + swa_branching_seqlen: Option, + #[pyo3(attribute)] mamba_host_hit_length: usize, #[pyo3(attribute)] mamba_branching_seqlen: Option, @@ -515,6 +517,7 @@ pub struct InsertParamsBinding { pub mamba_value: Option>, pub prev_prefix_len: usize, pub swa_evicted_seqlen: usize, + pub swa_branching_seqlen: Option, 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, prev_prefix_len: usize, swa_evicted_seqlen: usize, + swa_branching_seqlen: Option, chunked: bool, priority: i64, mamba_value: Option>, @@ -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, mamba_host_hit_length: usize, mamba_branching_seqlen: Option, 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, host_insert_dropped: bool, mamba_exist: bool, + swa_branch_inserted: bool, adopted_ranges: Option>>, cache_actions: Py, } @@ -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 TreeCoreBinding { 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 TreeCoreBinding { 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 TreeCoreBinding { .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 TreeCoreBinding { 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 TreeCoreBinding { 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) diff --git a/rust/sglang-radix-tree/src/tests/components/base.rs b/rust/sglang-radix-tree/src/tests/components/base.rs index 6554cd40d..270fa5d71 100644 --- a/rust/sglang-radix-tree/src/tests/components/base.rs +++ b/rust/sglang-radix-tree/src/tests/components/base.rs @@ -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, diff --git a/rust/sglang-radix-tree/src/tests/components/full.rs b/rust/sglang-radix-tree/src/tests/components/full.rs index dd0725123..18ad7dae3 100644 --- a/rust/sglang-radix-tree/src/tests/components/full.rs +++ b/rust/sglang-radix-tree/src/tests/components/full.rs @@ -72,6 +72,7 @@ fn insert(tc: &mut UnifiedTreeCore>, key: &Vec, 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, diff --git a/rust/sglang-radix-tree/src/tests/components/mamba.rs b/rust/sglang-radix-tree/src/tests/components/mamba.rs index ac9c586bc..672f70bac 100644 --- a/rust/sglang-radix-tree/src/tests/components/mamba.rs +++ b/rust/sglang-radix-tree/src/tests/components/mamba.rs @@ -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); diff --git a/rust/sglang-radix-tree/src/tests/components/swa.rs b/rust/sglang-radix-tree/src/tests/components/swa.rs index 9e73ce969..8f53a7236 100644 --- a/rust/sglang-radix-tree/src/tests/components/swa.rs +++ b/rust/sglang-radix-tree/src/tests/components/swa.rs @@ -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> { + 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>, node: NodeIdx_, value: i64) { + tc.arena + .set_device_value(node, SWA, Tensor::from_slice(&[value])); +} + +fn backup_transfers( + tc: &UnifiedTreeCore>, + window: usize, + node: NodeIdx_, +) -> Option> { + 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>, + window: usize, + node: NodeIdx_, +) -> (Vec, Vec) { + 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::::try_from(device_indices).unwrap(), + ) +} + +fn commit_backup( + tc: &mut UnifiedTreeCore>, + node: NodeIdx_, + host_indices: &[i64], + nodes_to_load: Option>, +) { + 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::>::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::>::needs_incremental_backup( + &swa, &tc, b + )); + + set_swa_host(&mut tc, a); + assert!(!TreeComponent::>::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::>::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::>::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::>::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> = 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>, + device_len: usize, + host_hit_length: usize, + full_kv_hit_length: usize, +) -> Option { + 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:?}" + ); + } +} diff --git a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs index cb8302c70..c9549ead2 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs @@ -1863,6 +1863,7 @@ fn insert_params<'k>(key: &'k Vec, value: &[i64]) -> InsertParams<'k, Vec( mamba_value, prev_prefix_len, swa_evicted_seqlen: 0, + swa_branching_seqlen: None, chunked: false, priority: 0, track_adopted_ranges: false, diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index 185054280..871075295 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -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, /// 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, /// The donated mamba slot for the insert target leaf; None on non-mamba trees. pub mamba_value: Option, /// 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, /// Whether write-through rejected a host suffix below an unbacked parent. @@ -207,6 +213,7 @@ pub struct InsertWalkState { namespace: KeyNamespace, prev_prefix_len: usize, swa_evicted_seqlen: usize, + swa_branching_seqlen: Option, mamba_value: Option, chunked: bool, priority: i64, @@ -532,6 +539,9 @@ pub struct UnifiedTreeCore { 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 UnifiedTreeCore { 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 UnifiedTreeCore { 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 UnifiedTreeCore { 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 UnifiedTreeCore { 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 UnifiedTreeCore { 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 UnifiedTreeCore { 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 UnifiedTreeCore { 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 UnifiedTreeCore { /// 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 UnifiedTreeCore { 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 UnifiedTreeCore { 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 UnifiedTreeCore { 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 UnifiedTreeCore { inserted_host_node: None, host_insert_dropped: false, mamba_exist: false, + swa_branch_inserted: false, adopted_ranges: None, cache_actions, }; diff --git a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py index 6a41bd98c..688e8bbff 100644 --- a/test/registered/unit/mem_cache/test_rust_tree_core_integration.py +++ b/test/registered/unit/mem_cache/test_rust_tree_core_integration.py @@ -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__])) diff --git a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py index a67db3a67..d7f8e8b7d 100644 --- a/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py +++ b/test/registered/unit/mem_cache/test_unified_radix_cache_unittest.py @@ -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