[UnifiedTree] fix: drop prefetched host refill under an un-backed-up parent (#31902)

Co-authored-by: hzh0425 <hzh0425@apache.org>

python/sglang/srt/mem_cache/base_prefix_cache.py

@@ -83,6 +83,7 @@ class InsertResult:
     last_device_node: Any = None
     mamba_exist: bool = False
     inserted_host_node: Any = None
+    host_insert_dropped: bool = False
     # Controller-applied actions from the non-stepped channels (e.g. insert_host); the stepped insert emits via InsertStepResult.actions.
     cache_actions: list[CacheAction | ComponentAction] = dataclasses.field(
         default_factory=list

python/sglang/srt/mem_cache/unified_cache/unified_tree_core.py

@@ -1494,6 +1494,16 @@ class UnifiedTreeCore(UnifiedTreeCoreInterface):
                 result.inserted_host_node = node.id
             return result
 
+        # Drop the refill only under write-through (a non-write-back policy).
+        if node is not self.root_node and not node.backuped and not self.is_write_back:
+            logger.info(
+                "HiCache prefetch dropped %d-token refill under un-backed-up node %d",
+                len(host_value),
+                node.id,
+            )
+            result.host_insert_dropped = True
+            return result
+
         new_node = self._new_node(priority=node.priority)
         new_node.parent = node
         new_node.key = key

python/sglang/srt/mem_cache/unified_radix_cache.py

@@ -1281,39 +1281,51 @@ class UnifiedRadixCache(BasePrefixCache):
 
         # Apply the host-insert walk's actions before the transfer commit.
         self._apply_cache_actions(insert_result.cache_actions)
-        commit_actions: list[CacheAction | ComponentAction] = []
-        self.tree_core.commit_hicache_transfers(
-            last_host_node_id,
-            CacheTransferPhase.PREFETCH,
-            comp_xfers,
-            cache_actions=commit_actions,
-            insert_result=insert_result,
-            pool_storage_result=operation.pool_storage_result,
-        )
-        self._apply_cache_actions(commit_actions)
-        # The commit emits via commit_actions only; the walk's were applied above.
-        assert not insert_result.cache_actions
 
-        self.cache_controller.mem_pool_host.free(
-            host_indices[: insert_result.prefix_len]
-        )
-        self.cache_controller.append_host_mem_release(
-            host_indices[min_completed_tokens:completed_tokens]
-        )
+        if insert_result.host_insert_dropped:
+            self.cache_controller.append_host_mem_release(
+                host_indices=host_indices[:completed_tokens],
+                extra_pools=[x for xfers in comp_xfers.values() for x in xfers],
+            )
+            loaded_from_storage = 0
+            released_tokens = completed_tokens
+        else:
+            commit_actions: list[CacheAction | ComponentAction] = []
+            self.tree_core.commit_hicache_transfers(
+                last_host_node_id,
+                CacheTransferPhase.PREFETCH,
+                comp_xfers,
+                cache_actions=commit_actions,
+                insert_result=insert_result,
+                pool_storage_result=operation.pool_storage_result,
+            )
+            self._apply_cache_actions(commit_actions)
+            # The commit emits via commit_actions only; the walk's were applied above.
+            assert not insert_result.cache_actions
+
+            self.cache_controller.mem_pool_host.free(
+                host_indices[: insert_result.prefix_len]
+            )
+            self.cache_controller.append_host_mem_release(
+                host_indices[min_completed_tokens:completed_tokens]
+            )
+            loaded_from_storage = min_completed_tokens - insert_result.prefix_len
+            released_tokens = completed_tokens - min_completed_tokens
+
         self.dec_host_lock_ref(last_host_node_id, anchor_lock_params)
         del self.ongoing_prefetch[req_id]
         self.cache_controller.prefetch_tokens_occupied -= len(prefetch_key)
 
-        loaded_from_storage = min_completed_tokens - insert_result.prefix_len
         self.prefetch_loaded_tokens_by_reqid[req_id] = loaded_from_storage
         logger.info(
-            "HiCache prefetch success req=%s completed_local=%d completed_synced=%d matched=%d loaded=%d tail_release=%d occupied=%d",
+            "HiCache prefetch %s req=%s completed_local=%d completed_synced=%d matched=%d loaded=%d released=%d occupied=%d",
+            "dropped" if insert_result.host_insert_dropped else "success",
             req_id,
             completed_tokens,
             min_completed_tokens,
             insert_result.prefix_len,
             loaded_from_storage,
-            completed_tokens - min_completed_tokens,
+            released_tokens,
             self.cache_controller.prefetch_tokens_occupied,
         )
         if self.enable_storage_metrics and self.storage_metrics_collector is not None: