diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index 01a07ad53..af883a455 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -732,6 +732,13 @@ class EAGLEWorkerV2(BaseSpecWorker): model_worker_batch ) assert verify_input.is_verify_input() + # Record a CUDA event after draft() GPU work is dispatched. + # This event will be waited on by plan_stream in verify() + # to ensure draft CUDA graph kernels finish before plan_stream + # begins metadata preparation. + if self.plan_stream: + self._draft_done_event = torch.get_device_module(self.device).Event() + self._draft_done_event.record() model_worker_batch.spec_info = verify_input batch_output = self.verify(model_worker_batch) with self.draft_worker.draft_tp_context( @@ -758,6 +765,12 @@ class EAGLEWorkerV2(BaseSpecWorker): # Batch 1: Target verify # Prepare for target verify in a separate stream with self.plan_stream_ctx: + # Wait for the draft CUDA graph to finish before plan_stream + # begins its work. Using an event is more targeted than + # wait_stream(main_stream) — it only waits for draft GPU + # work, not all queued main_stream operations. + if self.plan_stream and hasattr(self, "_draft_done_event"): + self.plan_stream.wait_event(self._draft_done_event) verify_forward_batch, can_run_cuda_graph = ( verify_input.prepare_for_v2_verify( self.req_to_token_pool, diff --git a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py index 40eb2ad66..b54dc3363 100644 --- a/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py +++ b/python/sglang/srt/speculative/multi_layer_eagle_worker_v2.py @@ -684,6 +684,10 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): draft_input: EagleDraftInput = model_worker_batch.spec_info verify_input: EagleVerifyInput = self.draft_worker.draft(model_worker_batch) assert verify_input.is_verify_input() + # Record a CUDA event after draft() GPU work is dispatched. + if self.plan_stream: + self._draft_done_event = torch.get_device_module(self.device).Event() + self._draft_done_event.record() model_worker_batch.spec_info = verify_input batch_output = self.verify(model_worker_batch) self.draft_worker._draft_extend_for_decode(model_worker_batch, batch_output) @@ -707,6 +711,10 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker): # Batch 1: Target verify # Prepare for target verify in a separate stream with self.plan_stream_ctx: + # Wait for the draft CUDA graph to finish before plan_stream + # begins its work. + if self.plan_stream and hasattr(self, "_draft_done_event"): + self.plan_stream.wait_event(self._draft_done_event) verify_forward_batch, can_run_cuda_graph = ( verify_input.prepare_for_v2_verify( self.req_to_token_pool,