[Spec] Publish DFLASH verify read-done event for fine-grained WAR barrier (#29541)
This commit is contained in:
@@ -995,14 +995,13 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
|
|||||||
)
|
)
|
||||||
with timer_ctx, self.backend.replay_session():
|
with timer_ctx, self.backend.replay_session():
|
||||||
self.load_batch(forward_batch, pp_proxy_tensors)
|
self.load_batch(forward_batch, pp_proxy_tensors)
|
||||||
# Snapshot built -- publish a read-done event for the WAR barrier.
|
# Publish a read-done event for the WAR barrier: a cuda-graph forward
|
||||||
# Only plain DECODE: the captured decode graph reads only its static
|
# finishes its shared req_to_token / SWA reads at this pre-replay
|
||||||
# snapshot, so the forward is done reading the shared pool here. Spec
|
# snapshot, so plain DECODE and DFLASH TARGET_VERIFY both qualify.
|
||||||
# verify (target_verify/dllm_extend) replays on this runner too, but
|
if forward_batch.forward_mode.is_decode() or (
|
||||||
# those are NOT the step's last shared-buffer-reading phase (eagle
|
forward_batch.forward_mode.is_target_verify()
|
||||||
# publishes from draft_extend; ngram/dflash must not publish here),
|
and self.model_runner.spec_algorithm.is_dflash()
|
||||||
# and some verify graphs may read beyond the snapshot in replay.
|
):
|
||||||
if forward_batch.forward_mode.is_decode():
|
|
||||||
read_done = self.device_module.Event()
|
read_done = self.device_module.Event()
|
||||||
read_done.record()
|
read_done.record()
|
||||||
self.model_runner.war_fastpath_read_done_event = read_done
|
self.model_runner.war_fastpath_read_done_event = read_done
|
||||||
|
|||||||
Reference in New Issue
Block a user