[Fix] Chain the seq_lens publish event records so prebuilt seeding keeps the forward fence (#30435)
This commit is contained in:
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user