[Spec] Allow speculative workers to stage prefill shared reads (#38554)

python/sglang/srt/model_executor/model_runner.py

@@ -20,7 +20,7 @@ import inspect
 import logging
 import time
 from dataclasses import dataclass
-from typing import Optional, Union
+from typing import Callable, Optional, Union
 
 import torch
 import torch.distributed as dist
@@ -438,6 +438,9 @@ class ModelRunner:
         # Read-done mailbox: the scheduler's WAR barrier reads it from the runner
         # its worker names, and treats None as the coarse whole-forward fence.
         self.shared_read_done_event: Optional[torch.cuda.Event] = None
+        # Scoped by a speculative worker to stage its shared reads before
+        # the target prefill graph publishes the read-done event.
+        self.prefill_shared_read_stager: Optional[Callable[[ForwardBatch], bool]] = None
 
         # CPU offload
         set_offloader(create_offloader(dp_rank=self.ps.dp_rank))

python/sglang/srt/model_executor/runner_utils/shared_read_event.py

@@ -31,18 +31,18 @@ def maybe_publish_prefill_shared_read_done(
         return
     if forward_batch.forward_mode != ForwardMode.EXTEND:
         return
-    # TODO(Jialin): Relax for EAGLE/MTP after validating the later
-    # draft-extend reader's WAR boundary.
-    if (
-        not model_runner.spec_algorithm.is_none()
-        and not model_runner.spec_algorithm.is_dflash_family()
-    ):
-        # Other speculative algorithms may have a later draft-extend reader.
-        return
     # The record lands right after replay prep, so PRE_REPLAY only.
     declared = model_runner.attn_backend.shared_read_ends(forward_batch.forward_mode)
     if declared is not SharedReadEnds.PRE_REPLAY:
         return
+    if (
+        not model_runner.spec_algorithm.is_none()
+        and not model_runner.spec_algorithm.is_dflash_family()
+    ):
+        # Stage the draft's shared reads before publishing the read-done event.
+        stage = getattr(model_runner, "prefill_shared_read_stager", None)
+        if stage is None or not stage(forward_batch):
+            return
     logger.info_once(
         "Prefill shared-read-done fastpath active (%s)",
         type(model_runner.attn_backend).__name__,