diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 1b506b826c7e..79ed493c93fc 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/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) diff --git a/test/registered/unit/model_executor/test_full_prefill_graph_hicache_load_fence.py b/test/registered/unit/model_executor/test_full_prefill_graph_hicache_load_fence.py new file mode 100644 index 000000000000..b3315f390f70 --- /dev/null +++ b/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()