Revert "[AMD]fix: use CUDA event for targeted draft-to-verify sync in… (#25917)
This commit is contained in:
@@ -807,13 +807,6 @@ class EAGLEWorkerV2(BaseSpecWorker):
|
|||||||
):
|
):
|
||||||
verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
|
verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
|
||||||
assert verify_input.is_verify_input()
|
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()
|
|
||||||
batch.spec_info = verify_input
|
batch.spec_info = verify_input
|
||||||
batch_output = self.verify(batch)
|
batch_output = self.verify(batch)
|
||||||
with (
|
with (
|
||||||
@@ -979,12 +972,6 @@ class EAGLEWorkerV2(BaseSpecWorker):
|
|||||||
# Batch 1: Target verify
|
# Batch 1: Target verify
|
||||||
# Prepare for target verify in a separate stream
|
# Prepare for target verify in a separate stream
|
||||||
with self.plan_stream_ctx:
|
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_forward_batch, can_run_cuda_graph = (
|
||||||
verify_input.prepare_for_v2_verify(
|
verify_input.prepare_for_v2_verify(
|
||||||
self.req_to_token_pool,
|
self.req_to_token_pool,
|
||||||
|
|||||||
@@ -704,10 +704,6 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
|
|||||||
)
|
)
|
||||||
verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
|
verify_input: EagleVerifyInput = self.draft_worker.draft(batch)
|
||||||
assert verify_input.is_verify_input()
|
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()
|
|
||||||
batch.spec_info = verify_input
|
batch.spec_info = verify_input
|
||||||
batch_output = self.verify(batch)
|
batch_output = self.verify(batch)
|
||||||
self.draft_worker._draft_extend_for_decode(batch, batch_output)
|
self.draft_worker._draft_extend_for_decode(batch, batch_output)
|
||||||
@@ -726,10 +722,6 @@ class MultiLayerEagleWorkerV2(BaseSpecWorker):
|
|||||||
# Batch 1: Target verify
|
# Batch 1: Target verify
|
||||||
# Prepare for target verify in a separate stream
|
# Prepare for target verify in a separate stream
|
||||||
with self.plan_stream_ctx:
|
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_forward_batch, can_run_cuda_graph = (
|
||||||
verify_input.prepare_for_v2_verify(
|
verify_input.prepare_for_v2_verify(
|
||||||
self.req_to_token_pool,
|
self.req_to_token_pool,
|
||||||
|
|||||||
Reference in New Issue
Block a user