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 3ee02de12..796a39912 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 @@ -2116,6 +2116,7 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface): self._update_evictable_leaf_sets(new_node) self._update_evictable_leaf_sets(node) result.inserted_host_node = new_node.id + self.kv_events.record_store(new_node, medium=StorageMedium.CPU) return result def build_backup_spec(self, node_id: NodeId): 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 c9549ead2..01a8fb208 100644 --- a/rust/sglang-radix-tree/src/tests/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/tests/unified_tree_core.rs @@ -3634,6 +3634,40 @@ fn insert_host_attaches_a_host_only_leaf_under_the_root() { tc.sanity_check(&[], &[]); } +#[test] +fn insert_host_publishes_a_host_store_event() { + // A storage-prefetch refill has no write-through ack to publish it. + let mut tc = events_core(2); + let root = tc.arena.root(); + let key = vec![1i64, 2, 7, 8]; + let hashes = crate::node::get_hash_str::>(&key, None, 2); + let result = tc + .insert_host( + tc.arena.node(root).id, + /* extra_key = */ None, + key.clone(), + Tensor::from_slice(&[100i64, 101, 102, 103]), + hashes.clone(), + ) + .expect("live test node"); + assert!(!result.host_insert_dropped); + assert!(result.inserted_host_node.is_some()); + assert_eq!( + tc.take_events(), + vec![KvCacheEvent::BlockStored { + block_hashes: hashes + .iter() + .map(|hash| crate::node::hash_str_to_int64(hash)) + .collect(), + parent_block_hash: None, + token_ids: key, + block_size: 2, + medium: StorageMedium::Cpu, + cache_salt: None, + }] + ); +} + #[test] fn insert_host_allows_a_suffix_under_an_unbacked_write_back_parent() { let mut tc = core(); diff --git a/rust/sglang-radix-tree/src/unified_tree_core.rs b/rust/sglang-radix-tree/src/unified_tree_core.rs index 871075295..44ab0acff 100644 --- a/rust/sglang-radix-tree/src/unified_tree_core.rs +++ b/rust/sglang-radix-tree/src/unified_tree_core.rs @@ -3107,6 +3107,7 @@ impl UnifiedTreeCore { self.update_evictable_leaf_sets_(new_node_id); self.update_evictable_leaf_sets_(node_id); result.inserted_host_node = Some(self.arena.node(new_node_id).id); + self.record_store_event_(new_node_id, StorageMedium::Cpu); Ok(result) } 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 d7f8e8b7d..745074fa5 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 @@ -1208,6 +1208,36 @@ class TestUnifiedRadixCacheKVEvents(CustomTestCase): removed_cpu = self._removed_events(cache, StorageMedium.CPU) self.assertCountEqual(self._event_hashes(removed_cpu), stored_hashes) + def test_hicache_storage_prefetch_publishes_host_only_suffix(self): + """A storage-prefetch refill has no write-through ack to publish it.""" + from sglang.srt.mem_cache.utils import get_hash_str, hash_str_to_int64 + + cache, _, _ = build_fixture(self.cfg, enable_kv_cache_events=True) + self._init_hicache(cache) + cache.take_events() # Clear reset / init events. + + tokens = [1, 2, 7, 8] + hash_values = get_hash_str(array("q", tokens), None, page_size=cache.page_size) + result = cache.tree_core.insert_host( + cache.root_node_handle(), + RadixKey(array("q", tokens)), + torch.tensor([100, 101, 102, 103], dtype=torch.int64), + hash_values, + ) + self.assertFalse(result.host_insert_dropped) + self.assertIsNotNone(result.inserted_host_node) + + # The recorder coalesces the parent-linked pages into one event. + stored_cpu = self._stored_events(cache, StorageMedium.CPU) + self.assertEqual(len(stored_cpu), 1) + self.assertEqual(list(stored_cpu[0].token_ids), tokens) + self.assertEqual(stored_cpu[0].block_size, cache.page_size) + self.assertEqual( + self._event_hashes(stored_cpu), + [hash_str_to_int64(value) for value in hash_values], + ) + self.assertIsNone(stored_cpu[0].parent_block_hash) + def test_hicache_split_pending_write_through_publishes_fragments(self): cache, allocator, _ = build_fixture(self.cfg, enable_kv_cache_events=True) self._init_hicache(cache)