[HiCache] Publish a host store event for storage-prefetch refills (#38486)

Co-authored-by: Shuwen Wang <47200617+alphabetc1@users.noreply.github.com>

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):

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::<Vec<i64>>(&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();

rust/sglang-radix-tree/src/unified_tree_core.rs

@@ -3107,6 +3107,7 @@ impl<K: ChildKeyType> UnifiedTreeCore<K> {
         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)
     }
 

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)