[HiCache] Fence full prefill CUDA-graph replay behind pending load-back

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -76,6 +76,7 @@
 )
 from sglang.srt.layers.logits_processor import LogitsProcessorOutput
 from sglang.srt.layers.pooler import EmbeddingPoolerOutput
+from sglang.srt.managers.cache_controller import LayerDoneCounter
 from sglang.srt.model_executor.cuda_graph_buffer_registry import (
     CudaGraphBufferRegistry,
     build_prefill_registry,
@@ -2096,12 +2097,28 @@ def _validate_capture_hidden_mode(self, forward_batch: ForwardBatch) -> None:
                 f"capture mode ({self.capture_hidden_mode.name})."
             )
 
+    def _wait_for_hicache_load_back(self) -> None:
+        """Order the forward stream behind the pending HiCache load-back.
+
+        The pools' per-layer ``layer_transfer_counter.wait_until`` calls are
+        Python-side and only run on replay inside eager graph breaks. The full
+        backend replays the whole transformer body as one graph, so wait once
+        on the load op's final event before replay.
+        """
+        counter = self.model_runner.token_to_kv_pool.layer_transfer_counter
+        if isinstance(counter, LayerDoneCounter) and counter.consumer_index >= 0:
+            self.device_module.current_stream().wait_event(
+                counter.events[counter.consumer_index].finish_event
+            )
+
     def execute(
         self, forward_batch: ForwardBatch, **kwargs
     ) -> Union[LogitsProcessorOutput, PPProxyTensors, EmbeddingPoolerOutput]:
         self._validate_capture_hidden_mode(forward_batch)
         with self.backend.replay_session():
             static_forward_batch = self.load_batch(forward_batch, **kwargs)
+            if self._is_full_backend:
+                self._wait_for_hicache_load_back()
             static_num_tokens = len(static_forward_batch.input_ids)
             raw_num_tokens = self.raw_num_tokens
             shape_key = self._shape_key(static_num_tokens, forward_batch)

test/registered/unit/model_executor/test_full_prefill_graph_hicache_load_fence.py

@@ -0,0 +1,51 @@
+import unittest
+from types import SimpleNamespace
+
+from sglang.srt.managers.cache_controller import LayerDoneCounter
+from sglang.srt.model_executor.runner.prefill_cuda_graph_runner import (
+    PrefillCudaGraphRunner,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+from sglang.test.test_utils import CustomTestCase
+
+register_cpu_ci(est_time=5, suite="base-a-test-cpu")
+
+
+def _counter(consumer_index):
+    counter = object.__new__(LayerDoneCounter)
+    counter.events = [SimpleNamespace(finish_event=f"finish-{i}") for i in range(3)]
+    counter.producer_index = consumer_index
+    counter.consumer_index = consumer_index
+    return counter
+
+
+def _runner(counter):
+    waited = []
+    stream = SimpleNamespace(wait_event=waited.append)
+    runner = object.__new__(PrefillCudaGraphRunner)
+    runner.device_module = SimpleNamespace(current_stream=lambda: stream)
+    runner.model_runner = SimpleNamespace(
+        token_to_kv_pool=SimpleNamespace(layer_transfer_counter=counter)
+    )
+    return runner, waited
+
+
+class TestFullPrefillGraphHiCacheLoadFence(CustomTestCase):
+    def test_pending_load_back_waits_on_final_event(self):
+        runner, waited = _runner(_counter(consumer_index=1))
+        runner._wait_for_hicache_load_back()
+        self.assertEqual(waited, ["finish-1"])
+
+    def test_no_pending_load_back_no_wait(self):
+        runner, waited = _runner(_counter(consumer_index=-1))
+        runner._wait_for_hicache_load_back()
+        self.assertEqual(waited, [])
+
+    def test_no_hicache_no_wait(self):
+        runner, waited = _runner(None)
+        runner._wait_for_hicache_load_back()
+        self.assertEqual(waited, [])
+
+
+if __name__ == "__main__":
+    unittest.main()