[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__,

test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py

@@ -1,5 +1,5 @@
 from types import SimpleNamespace
-from unittest.mock import create_autospec
+from unittest.mock import Mock, call, create_autospec
 
 import pytest
 
@@ -35,6 +35,7 @@ def _model_runner(*, spec_algorithm=SpeculativeAlgorithm.NONE, compliant=True):
         spec_algorithm=spec_algorithm,
         attn_backend=attn_backend,
         shared_read_done_event=None,
+        prefill_shared_read_stager=None,
     )
 
 
@@ -65,10 +66,52 @@ def test_disabled_when_flag_is_false():
 )
 def test_dflash_family_target_prefill_publishes(algorithm):
     runner = _model_runner(spec_algorithm=algorithm)
+    runner.prefill_shared_read_stager = Mock(return_value=False)
     with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True):
         maybe_publish_prefill_shared_read_done(runner, _batch(), _DEVICE_MODULE)
     published = runner.shared_read_done_event
     assert isinstance(published, _Event) and published.recorded
+    runner.prefill_shared_read_stager.assert_not_called()
+
+
+@pytest.mark.parametrize("staged", [False, True])
+def test_speculative_prefill_publishes_only_after_staging(staged):
+    runner, batch = _model_runner(spec_algorithm=SpeculativeAlgorithm.EAGLE), _batch()
+    calls = Mock()
+    runner.prefill_shared_read_stager = calls.stage
+    calls.stage.return_value = staged
+    calls.Event.side_effect = _Event
+    with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True):
+        maybe_publish_prefill_shared_read_done(
+            runner, batch, SimpleNamespace(Event=calls.Event)
+        )
+    assert calls.mock_calls == [call.stage(batch)] + ([call.Event()] if staged else [])
+    published = runner.shared_read_done_event
+    if staged:
+        assert isinstance(published, _Event) and published.recorded
+    else:
+        assert published is None
+
+
+@pytest.mark.parametrize(
+    "enabled,mode,compliant",
+    [
+        (False, ForwardMode.EXTEND, True),
+        (True, ForwardMode.TARGET_VERIFY, True),
+        (True, ForwardMode.MIXED, True),
+        (True, ForwardMode.DECODE, True),
+        (True, ForwardMode.EXTEND, False),
+    ],
+)
+def test_prefill_gates_skip_staging(enabled, mode, compliant):
+    runner = _model_runner(
+        spec_algorithm=SpeculativeAlgorithm.EAGLE, compliant=compliant
+    )
+    runner.prefill_shared_read_stager = Mock(return_value=True)
+    with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(enabled):
+        maybe_publish_prefill_shared_read_done(runner, _batch(mode), _DEVICE_MODULE)
+    runner.prefill_shared_read_stager.assert_not_called()
+    assert runner.shared_read_done_event is None
 
 
 def test_gates_exclude_non_prefill_unsupported_algorithm_and_noncompliant_backend():