From da2f434951e0e4ea72b227078ee09c46f19ec2c8 Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Fri, 18 Sep 2026 13:11:01 -0700 Subject: [PATCH] [Spec] Add explicit prefill shared-read capability for plugins (#39502) --- .../runner_utils/shared_read_event.py | 5 +- python/sglang/srt/speculative/spec_info.py | 8 ++ .../sglang/srt/speculative/spec_registry.py | 4 + .../runner/test_prefill_shared_read_done.py | 94 ++++++++++++++++++- 4 files changed, 105 insertions(+), 6 deletions(-) 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 b0a29e6a3..8bb7c7c91 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 @@ -35,10 +35,7 @@ def maybe_publish_prefill_shared_read_done( 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() - ): + if not model_runner.spec_algorithm.supports_prefill_shared_read_done(): # 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): diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index f296e1fe3..1cb443feb 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -136,6 +136,14 @@ class SpeculativeAlgorithm(Enum): def supports_target_verify_for_draft(self) -> bool: return self.is_dflash_family() + def supports_prefill_shared_read_done(self) -> bool: + """Whether target EXTEND has no later speculative shared-buffer reader. + + The backend must still declare a pre-replay read end. Other algorithms + must stage their draft's shared reads before publishing the target event. + """ + return self.is_none() or self.is_dflash_family() + def supports_mixed_chunk(self) -> bool: """Whether mixed chunk prefill may stay enabled with this algorithm. diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index f7e486712..edef1d000 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/python/sglang/srt/speculative/spec_registry.py @@ -100,6 +100,10 @@ class CustomSpecAlgo: def supports_target_verify_for_draft(self) -> bool: return False + def supports_prefill_shared_read_done(self) -> bool: + # Whether target EXTEND has no later speculative shared-buffer reader. + return False + def supports_ragged_verify(self) -> bool: return False 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 e67a27b76..0c96e270a 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 @@ -13,7 +13,9 @@ from sglang.srt.model_executor.runner_utils import ( maybe_publish_prefill_shared_read_done, ) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm +from sglang.srt.speculative.spec_registry import CustomSpecAlgo from sglang.test.ci.ci_register import register_cpu_ci +from sglang.test.test_utils import CustomTestCase register_cpu_ci(est_time=11, suite="base-a-test-cpu") @@ -75,8 +77,20 @@ def test_dflash_family_target_prefill_publishes(algorithm): @pytest.mark.parametrize("staged", [False, True]) -def test_speculative_prefill_publishes_only_after_staging(staged): - runner, batch = _model_runner(spec_algorithm=SpeculativeAlgorithm.EAGLE), _batch() +@pytest.mark.parametrize( + "algorithm", + [ + SpeculativeAlgorithm.EAGLE, + SpeculativeAlgorithm.EAGLE3, + SpeculativeAlgorithm.FROZEN_KV_MTP, + SpeculativeAlgorithm.STANDALONE, + SpeculativeAlgorithm.NGRAM, + SpeculativeAlgorithm.UNO, + CustomSpecAlgo("TEST_PREFILL", lambda _: object), + ], +) +def test_speculative_prefill_publishes_only_after_staging(algorithm, staged): + runner, batch = _model_runner(spec_algorithm=algorithm), _batch() calls = Mock() runner.prefill_shared_read_stager = calls.stage calls.stage.return_value = staged @@ -123,6 +137,12 @@ def test_gates_exclude_non_prefill_unsupported_algorithm_and_noncompliant_backen (_model_runner(), _batch(ForwardMode.DECODE)), # The algorithm has a later prefill reader or unverified ownership. (_model_runner(spec_algorithm=SpeculativeAlgorithm.EAGLE), _batch()), + ( + _model_runner( + spec_algorithm=CustomSpecAlgo("TEST_PREFILL", lambda _: object) + ), + _batch(), + ), # Backend has not declared a pre-replay prefill read end. (_model_runner(compliant=False), _batch()), ): @@ -130,5 +150,75 @@ def test_gates_exclude_non_prefill_unsupported_algorithm_and_noncompliant_backen assert runner.shared_read_done_event is None +class _TargetOnlyPrefillAlgo(CustomSpecAlgo): + def supports_prefill_shared_read_done(self) -> bool: + return True + + +class TestPrefillReadDoneCapability(CustomTestCase): + def test_builtin_allowlist_is_unchanged(self): + allowed = { + SpeculativeAlgorithm.NONE, + SpeculativeAlgorithm.DFLASH, + SpeculativeAlgorithm.DSPARK, + } + for algorithm in SpeculativeAlgorithm: + with self.subTest(algorithm=algorithm): + runner = _model_runner(spec_algorithm=algorithm) + with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True): + maybe_publish_prefill_shared_read_done( + runner, _batch(), _DEVICE_MODULE + ) + self.assertEqual( + runner.shared_read_done_event is not None, algorithm in allowed + ) + + def test_custom_algorithm_requires_explicit_opt_in(self): + for algorithm_type, expected in ( + (CustomSpecAlgo, False), + (_TargetOnlyPrefillAlgo, True), + ): + with self.subTest(algorithm_type=algorithm_type): + algorithm = algorithm_type( + "TEST_PREFILL", lambda _: object, supports_overlap=True + ) + runner, batch = _model_runner(spec_algorithm=algorithm), _batch() + 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 + ) + event = runner.shared_read_done_event + self.assertEqual(event is not None, expected) + if expected: + self.assertTrue(event.recorded) + runner.prefill_shared_read_stager.assert_not_called() + else: + runner.prefill_shared_read_stager.assert_called_once_with(batch) + + def test_opt_in_does_not_bypass_mode_backend_or_flag(self): + algorithm = _TargetOnlyPrefillAlgo("TEST_PREFILL", lambda _: object) + cases = [(False, ForwardMode.EXTEND, SharedReadEnds.PRE_REPLAY)] + cases.extend( + (True, mode, SharedReadEnds.PRE_REPLAY) + for mode in ForwardMode + if mode != ForwardMode.EXTEND + ) + cases.extend( + (True, ForwardMode.EXTEND, declared) + for declared in SharedReadEnds + if declared != SharedReadEnds.PRE_REPLAY + ) + for enabled, mode, declared in cases: + with self.subTest(enabled=enabled, mode=mode, declared=declared): + runner = _model_runner(spec_algorithm=algorithm) + runner.attn_backend.shared_read_ends.return_value = declared + with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(enabled): + maybe_publish_prefill_shared_read_done( + runner, _batch(mode), _DEVICE_MODULE + ) + self.assertIsNone(runner.shared_read_done_event) + + if __name__ == "__main__": raise SystemExit(pytest.main([__file__, "-v"]))