diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 79a8aa859..4df5b8a7e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1488,12 +1488,8 @@ class Scheduler( if self.device == "cpu": self.schedule_stream.synchronize = lambda: None # No-op for CPU # The global WAR barrier fences the scheduler's next shared-buffer write - # on the forward's read-done event. DFLASH opts out: it fences its own - # req_to_token writes with verify_done / plan-stream deps, so the global - # barrier would only serialize plan overlap without adding correctness. - self._war_barrier_enabled = ( - is_cuda() or envs.SGLANG_ENABLE_WAR_BARRIER.get() - ) and not self.spec_algorithm.is_dflash() + # on the previous forward's read of the shared pool. + self._war_barrier_enabled = is_cuda() or envs.SGLANG_ENABLE_WAR_BARRIER.get() with self.device_module.StreamContext(self.schedule_stream): dispatch_event_loop(self) diff --git a/python/sglang/srt/speculative/dflash_info_v2.py b/python/sglang/srt/speculative/dflash_info_v2.py index 36177b459..fa4d35942 100644 --- a/python/sglang/srt/speculative/dflash_info_v2.py +++ b/python/sglang/srt/speculative/dflash_info_v2.py @@ -46,7 +46,6 @@ class DFlashDraftInputV2(SpecInput): bonus_tokens: torch.Tensor new_seq_lens: torch.Tensor hidden_states: torch.Tensor - verify_done: Optional[torch.cuda.Event] = None max_top_k: int = 1 uniform_top_k_value: Optional[int] = None reserved_seq_lens_cpu: Optional[torch.Tensor] = None @@ -113,7 +112,6 @@ class DFlashDraftInputV2(SpecInput): bonus_tokens=torch.empty((0,), device=device, dtype=torch.int64), new_seq_lens=torch.empty((0,), device=device, dtype=torch.int64), hidden_states=torch.empty((0, 0), device=device, dtype=torch.float16), - verify_done=None, ) def prepare_for_decode(self, batch: ScheduleBatch): @@ -126,9 +124,6 @@ class DFlashDraftInputV2(SpecInput): the overallocated mapping. """ plan_stream, plan_stream_ctx = _get_overlap_plan_stream(batch.device) - if plan_stream is None and self.verify_done is not None: - # Ensure previous forward is completed before mutating shared buffers. - self.verify_done.synchronize() bs = batch.batch_size() if bs == 0: @@ -192,9 +187,6 @@ class DFlashDraftInputV2(SpecInput): # The plan stream must wait for those writes before reading them. plan_stream.wait_stream(caller_stream) - if plan_stream is not None and self.verify_done is not None: - plan_stream.wait_event(self.verify_done) - cur_kv_lens = self._prepare_cur_kv_lens_gpu_buf[:bs] nxt_kv_lens = self._prepare_nxt_kv_lens_gpu_buf[:bs] cur_kv_lens.copy_(cur_kv_lens_cpu_t, non_blocking=True) diff --git a/python/sglang/srt/speculative/dflash_worker_v2.py b/python/sglang/srt/speculative/dflash_worker_v2.py index 94194665b..2553a2927 100644 --- a/python/sglang/srt/speculative/dflash_worker_v2.py +++ b/python/sglang/srt/speculative/dflash_worker_v2.py @@ -1233,7 +1233,6 @@ class DFlashWorkerV2(BaseSpecWorker): *, bonus_tokens: torch.Tensor, seq_lens: torch.Tensor, - verify_done: Optional[torch.cuda.Event] = None, ) -> DFlashDraftInputV2: bs = int(seq_lens.numel()) device = bonus_tokens.device @@ -1243,7 +1242,6 @@ class DFlashWorkerV2(BaseSpecWorker): bonus_tokens=bonus_tokens.to(dtype=torch.int64), new_seq_lens=seq_lens.to(dtype=torch.int64), hidden_states=torch.empty((bs, 0), device=device, dtype=torch.float16), - verify_done=verify_done, ) def _make_next_draft_input_decode( @@ -1251,7 +1249,6 @@ class DFlashWorkerV2(BaseSpecWorker): *, bonus_tokens: torch.Tensor, new_seq_lens: torch.Tensor, - verify_done: Optional[torch.cuda.Event] = None, ) -> DFlashDraftInputV2: bs = int(new_seq_lens.numel()) device = bonus_tokens.device @@ -1261,7 +1258,6 @@ class DFlashWorkerV2(BaseSpecWorker): bonus_tokens=bonus_tokens.to(dtype=torch.int64), new_seq_lens=new_seq_lens.to(dtype=torch.int64), hidden_states=torch.empty((bs, 0), device=device, dtype=torch.float16), - verify_done=verify_done, ) def forward_batch_generation( @@ -1331,9 +1327,6 @@ class DFlashWorkerV2(BaseSpecWorker): bonus_tokens=next_token_ids, seq_lens=batch.seq_lens, ) - verify_done = torch.get_device_module(device).Event() - verify_done.record() - batch_output.next_draft_input.verify_done = verify_done return batch_output # Decode / target-verify stage. @@ -1355,9 +1348,6 @@ class DFlashWorkerV2(BaseSpecWorker): ) if on_publish is not None: on_publish(next_draft_input.new_seq_lens) - verify_done = torch.get_device_module(self.device).Event() - verify_done.record() - next_draft_input.verify_done = verify_done return GenerationBatchResult( logits_output=None, next_token_ids=empty_ids, @@ -1732,9 +1722,6 @@ class DFlashWorkerV2(BaseSpecWorker): bonus_tokens=bonus, new_seq_lens=new_seq_lens, ) - verify_done = torch.get_device_module(device).Event() - verify_done.record() - next_draft_input.verify_done = verify_done return GenerationBatchResult( logits_output=logits_output,