[HiCache] Fence load-back behind the forward stream (#36738)

python/sglang/srt/managers/cache_controller.py

@@ -361,6 +361,9 @@ class HiCacheController:
         self.load_queue: List[CacheOperation] = []
         self.write_queue: List[CacheOperation] = []
         self.ack_load_queue: List[HiCacheAck] = []
+        # Set by the scheduler to the forward stream; gates load-back H2D
+        # behind in-flight forwards (see start_loading).
+        self.load_fence_stream = None
         self.ack_write_queue: List[HiCacheAck] = []
 
         self.l2_transfer_engine = L2TransferEngine(io_backend)
@@ -922,6 +925,14 @@ class HiCacheController:
         producer_event = self.layer_done_counter.events[producer_id]
         producer_event.start_event.record()
 
+        if self.load_fence_stream is not None:
+            # in overlap scheduling, reclaimed pages might still be written by the forward thread
+            # therefore a fence is needed for loading thread to prevent memory corruption
+            # todo: it's possible to use a finer-grained fence
+            self.l2_transfer_engine.host_to_device_stream.wait_stream(
+                self.load_fence_stream
+            )
+
         completion = self.l2_transfer_engine.submit_host_to_device(
             self._l2_load_transfers(host_indices, device_indices, pool_transfers),
             start_event=producer_event.start_event,

python/sglang/srt/managers/scheduler.py

@@ -577,6 +577,12 @@ class Scheduler(
         self.token_to_kv_pool_allocator = result.token_to_kv_pool_allocator
         self.disable_radix_cache = result.disable_radix_cache
         self.tree_cache = result.tree_cache
+        if self.enable_hierarchical_cache:
+            cache_controller = self.tree_cache.cache_controller
+            if cache_controller is not None:
+                cache_controller.load_fence_stream = (
+                    self.tp_worker.model_runner.forward_stream
+                )
         self.emit_metrics_constants()
         self.maybe_init_hccl_dp_prewarm()