[Fix] Chain the seq_lens publish event records so prebuilt seeding keeps the forward fence (#30435)

This commit is contained in:
Liangsheng Yin
2026-07-08 13:43:09 -07:00
committed by GitHub
parent 47a6dfd708
commit 10e7f2925f
2 changed files with 9 additions and 1 deletions
+7 -1
View File
@@ -322,8 +322,14 @@ class FutureMap:
self.new_seq_lens_buf[indices] = new_seq_lens.to(self.new_seq_lens_buf.dtype)
# Only spec_v2 needs the event; it gates the seq_lens D2H on the private stream.
if self.spec_algo.is_some():
device_module = torch.get_device_module(self.device)
if self.publish_ready is None:
self.publish_ready = torch.get_device_module(self.device).Event()
self.publish_ready = device_module.Event()
else:
# Chain the records: event fire implies every prior publish is
# visible, so an off-forward-stream publish (PD-decode prebuilt
# seeding) cannot drop the in-flight forward's fence.
device_module.current_stream().wait_event(self.publish_ready)
self.publish_ready.record()
def stash(self, future_indices: torch.Tensor, payload: RelayPayload) -> None:
@@ -61,6 +61,8 @@ def build_eagle_disagg_draft_input(
if batch.enable_overlap:
spec_info.future_indices = batch.req_pool_indices
# Seed the relay buf with the known seq_lens; publish's chained record
# keeps the in-flight forward's fence intact (see FutureMap.publish).
future_map.publish(spec_info.future_indices, batch.seq_lens)
future_map.stash(
spec_info.future_indices, RelayPayload.from_draft_input(spec_info)