diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 14ad4fb4a..28d1beae5 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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)) diff --git a/python/sglang/srt/model_executor/runner_utils/shared_read_event.py b/python/sglang/srt/model_executor/runner_utils/shared_read_event.py index 4d9f36bb6..b0a29e6a3 100644 --- a/python/sglang/srt/model_executor/runner_utils/shared_read_event.py +++ b/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__, diff --git a/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py b/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py index 2eb82a753..e67a27b76 100644 --- a/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py +++ b/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():