From 711bdacb825231d01d8d0da2431d2ccdc0b28c1f Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Mon, 17 Aug 2026 01:35:29 -0700 Subject: [PATCH] [Spec] Resolve shared-read ends from the backend declaration alone (#35059) --- python/sglang/srt/managers/scheduler.py | 8 +--- .../sglang/srt/model_executor/model_runner.py | 5 +-- .../runner/decode_cuda_graph_runner.py | 10 ----- python/sglang/srt/speculative/ngram_worker.py | 2 + python/sglang/srt/speculative/spec_info.py | 8 ---- .../sglang/srt/speculative/spec_registry.py | 4 -- ...est_decode_cuda_graph_shared_read_fence.py | 37 +++++++------------ 7 files changed, 20 insertions(+), 54 deletions(-) diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 58f9e4eb6..4bac671ad 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1701,12 +1701,8 @@ class Scheduler( dispatch_event_loop(self) def _apply_war_barrier(self): - # Called right after each launch: order later schedule_stream work - # (result processing, next iteration's writes) behind the forward's - # shared-buffer reads. Fast path: wait on the read-done event the - # forward published after its snapshot (non-spec: decode graph; spec: - # draft_extend), then clear it. Else whole-forward wait_stream - # (forceable via SGLANG_FORCE_COARSE_WAR_BARRIER). + # WAR: keep later schedule_stream writes behind this forward's shared reads. + # Clearing matters: a phase that skips the publish then falls back to coarse. if not self._war_barrier_enabled: return runner = self.model_worker.last_shared_read_runner diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 075c29371..9bcfd8085 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -403,9 +403,8 @@ class ModelRunner: # Init forward stream for overlap schedule self.forward_stream = torch.get_device_module(self.device).Stream() - # Published by the step's last shared-buffer-reading phase (decode graph, - # eagle draft extend, or prefill); the scheduler's WAR barrier waits on it - # then clears it. None -> coarse whole-forward wait_stream. + # 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 # CPU offload diff --git a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py index 01babb0cf..835c85899 100644 --- a/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py @@ -486,17 +486,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): return self.attn_backend def _resolve_shared_read_ends(self, attn_backend, forward_mode) -> SharedReadEnds: - """The backend's declaration, demoted when this runner cannot record - there. UNKNOWN records nothing (scheduler keeps the coarse fence).""" - if forward_mode.is_target_verify(): - if not self.model_runner.spec_algorithm.is_last_shared_read_phase( - forward_mode - ): - return SharedReadEnds.UNKNOWN - elif not forward_mode.is_decode(): - return SharedReadEnds.UNKNOWN declared = attn_backend.shared_read_ends(forward_mode) - if ( declared is SharedReadEnds.IN_REPLAY and self.in_graph_metadata_prep_done is None diff --git a/python/sglang/srt/speculative/ngram_worker.py b/python/sglang/srt/speculative/ngram_worker.py index 39caaa791..9bb6d3d04 100644 --- a/python/sglang/srt/speculative/ngram_worker.py +++ b/python/sglang/srt/speculative/ngram_worker.py @@ -419,6 +419,8 @@ class NGRAMWorker(BaseSpecWorker): batch_result = self.target_worker.forward_batch_generation( batch, is_verify=True ) + # Verify reads shared state past the in-graph marker; keep it coarse. + self.target_worker.model_runner.shared_read_done_event = None logits_output, can_run_cuda_graph = ( batch_result.logits_output, diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index e018ee4c0..440bcc0be 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -127,14 +127,6 @@ class SpeculativeAlgorithm(Enum): def supports_target_verify_for_draft(self) -> bool: return self.is_dflash_family() - def is_last_shared_read_phase(self, forward_mode) -> bool: - # The step's last shared-buffer-reading phase owns the shared-read-done - # publish. DSPARK has no draft_extend: its draft samples inside the - # verify graph, so verify is that last phase. - if self.is_dflash_family() or self.is_dspark(): - return forward_mode.is_target_verify() - return forward_mode.is_draft_extend_v2() - def supports_ragged_verify(self) -> bool: """Whether this algorithm's verify step may carry a RaggedVerifyLayout (per-request verify lengths); gates the token-bucket-keyed verify diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index a736f482a..4ab16a2a3 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/python/sglang/srt/speculative/spec_registry.py @@ -92,10 +92,6 @@ class CustomSpecAlgo: def supports_target_verify_for_draft(self) -> bool: return False - def is_last_shared_read_phase(self, forward_mode) -> bool: - # The step's last shared-buffer-reading phase owns the shared-read-done publish. - return forward_mode.is_draft_extend_v2() - def supports_ragged_verify(self) -> bool: return False diff --git a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_shared_read_fence.py b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_shared_read_fence.py index 5dab628cd..1c5dee227 100644 --- a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_shared_read_fence.py +++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_shared_read_fence.py @@ -16,18 +16,11 @@ from sglang.test.ci.ci_register import register_cpu_ci register_cpu_ci(est_time=1, suite="base-a-test-cpu") DECODE = ForwardMode.DECODE -VERIFY = ForwardMode.TARGET_VERIFY -EXTEND = ForwardMode.EXTEND -def _runner(*, owns_verify: bool = False, has_marker: bool = False): +def _runner(*, has_marker: bool = False): runner = DecodeCudaGraphRunner.__new__(DecodeCudaGraphRunner) - runner.model_runner = SimpleNamespace( - spec_algorithm=SimpleNamespace( - is_last_shared_read_phase=lambda fm: owns_verify and fm.is_target_verify() - ), - shared_read_done_event=None, - ) + runner.model_runner = SimpleNamespace(shared_read_done_event=None) runner.in_graph_metadata_prep_done = object() if has_marker else None return runner @@ -40,24 +33,22 @@ def _backend(declared: SharedReadEnds): @pytest.mark.parametrize( - "mode, owns_verify, declared, has_marker, expected", + "declared, has_marker, expected", [ - # Only decode / target verify publish; anything else keeps the coarse fence. - (EXTEND, False, SharedReadEnds.IN_REPLAY, True, SharedReadEnds.UNKNOWN), - # Target verify publishes only when it is the step's last reading phase. - (VERIFY, False, SharedReadEnds.IN_REPLAY, True, SharedReadEnds.UNKNOWN), - (VERIFY, True, SharedReadEnds.IN_REPLAY, True, SharedReadEnds.IN_REPLAY), - # A backend that keeps reading through the graph is never advanced. - (VERIFY, True, SharedReadEnds.POST_REPLAY, True, SharedReadEnds.POST_REPLAY), - # Nothing to demote: the declaration is honored as-is. - (DECODE, False, SharedReadEnds.IN_REPLAY, True, SharedReadEnds.IN_REPLAY), + # The backend's declaration decides where the record lands. + (SharedReadEnds.IN_REPLAY, True, SharedReadEnds.IN_REPLAY), # Nowhere to record in-graph -> fall back to the pre-replay record. - (DECODE, False, SharedReadEnds.IN_REPLAY, False, SharedReadEnds.PRE_REPLAY), + (SharedReadEnds.IN_REPLAY, False, SharedReadEnds.PRE_REPLAY), + # Only an in-graph declaration is demoted; the rest pass through. + (SharedReadEnds.POST_REPLAY, False, SharedReadEnds.POST_REPLAY), ], ) -def test_resolve_shared_read_ends(mode, owns_verify, declared, has_marker, expected): - runner = _runner(owns_verify=owns_verify, has_marker=has_marker) - assert runner._resolve_shared_read_ends(_backend(declared), mode) is expected +def test_resolve_shared_read_ends(declared, has_marker, expected): + runner = _runner(has_marker=has_marker) + backend = _backend(declared) + + assert runner._resolve_shared_read_ends(backend, DECODE) is expected + backend.shared_read_ends.assert_called_once_with(DECODE) def test_publish_read_done():