dflash: drop verify_done barrier; rely on scheduler WAR fallback (#29556)
This commit is contained in:
@@ -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)
|
||||
|
||||
|
||||
@@ -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)
|
||||
|
||||
@@ -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,
|
||||
|
||||
Reference in New Issue
Block a user