[Scheduler] Unify WAR read-done gating behind shared-read boundary declarations (#34052)

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

@@ -4,12 +4,12 @@ from types import SimpleNamespace
 import pytest
 import torch
 
+from sglang.srt.layers.attention.base_attn_backend import SharedReadBoundary
 from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors
 from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
     DecodeCudaGraphRunner,
 )
 from sglang.srt.model_executor.runner.shape_key import ShapeKey
-from sglang.srt.model_executor.runner_utils import WarReadDonePolicy
 from sglang.test.ci.ci_register import register_cpu_ci
 
 register_cpu_ci(est_time=1, suite="base-a-test-cpu")
@@ -19,14 +19,19 @@ class _SpecAlgorithm:
     def __init__(self, target_verify_war: bool = False):
         self._target_verify_war = target_verify_war
 
-    def supports_target_verify_war_read_done(self) -> bool:
-        return self._target_verify_war
+    def is_war_publish_phase(self, forward_mode) -> bool:
+        return self._target_verify_war and forward_mode.is_target_verify()
 
 
 def _attn_backend(*, breakable_metadata=False):
-    return SimpleNamespace(
-        use_captured_forward_metadata_for_breakable_cuda_graph=breakable_metadata,
-    )
+    def shared_read_boundary(forward_mode):
+        if breakable_metadata and forward_mode.is_target_verify():
+            return SharedReadBoundary.POST_REPLAY
+        if forward_mode.is_decode() or forward_mode.is_target_verify():
+            return SharedReadBoundary.IN_REPLAY
+        return SharedReadBoundary.UNKNOWN
+
+    return SimpleNamespace(shared_read_boundary=shared_read_boundary)
 
 
 def _runner(*, target_verify_war: bool = False, planted: bool = False):
@@ -42,35 +47,35 @@ def _runner(*, target_verify_war: bool = False, planted: bool = False):
     return runner
 
 
-def test_war_read_done_policy():
+def test_war_read_done_record():
     # Planted node: the graph re-arms it every replay.
     assert (
-        _runner(planted=True)._war_read_done_policy(_attn_backend(), ForwardMode.DECODE)
-        is WarReadDonePolicy.IN_GRAPH
+        _runner(planted=True)._war_read_done_record(_attn_backend(), ForwardMode.DECODE)
+        is SharedReadBoundary.IN_REPLAY
     )
-    # No node, snapshot backend: all shared reads finish before launch.
+    # No planted node: fall back to a pre-replay record.
     assert (
-        _runner()._war_read_done_policy(_attn_backend(), ForwardMode.DECODE)
-        is WarReadDonePolicy.PRE_REPLAY
+        _runner()._war_read_done_record(_attn_backend(), ForwardMode.DECODE)
+        is SharedReadBoundary.PRE_REPLAY
     )
     # Unrelated modes never publish from the decode graph runner.
     assert (
-        _runner(planted=True)._war_read_done_policy(_attn_backend(), ForwardMode.EXTEND)
-        is WarReadDonePolicy.NONE
+        _runner(planted=True)._war_read_done_record(_attn_backend(), ForwardMode.EXTEND)
+        is SharedReadBoundary.UNKNOWN
     )
-    # Backend placement cannot opt an unsupported algorithm into publication.
+    # The algorithm gate precedes the backend declaration.
     assert (
-        _runner(planted=True)._war_read_done_policy(
+        _runner(planted=True)._war_read_done_record(
             _attn_backend(breakable_metadata=True), ForwardMode.TARGET_VERIFY
         )
-        is WarReadDonePolicy.NONE
+        is SharedReadBoundary.UNKNOWN
     )
     # Captured-metadata verify keeps reading throughout the graph, even planted.
     assert (
-        _runner(target_verify_war=True, planted=True)._war_read_done_policy(
+        _runner(target_verify_war=True, planted=True)._war_read_done_record(
             _attn_backend(breakable_metadata=True), ForwardMode.TARGET_VERIFY
         )
-        is WarReadDonePolicy.POST_REPLAY
+        is SharedReadBoundary.POST_REPLAY
     )
 
 

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

@@ -3,6 +3,7 @@ from types import SimpleNamespace
 import pytest
 
 from sglang.srt.environ import envs
+from sglang.srt.layers.attention.base_attn_backend import SharedReadBoundary
 from sglang.srt.model_executor.forward_batch_info import ForwardMode
 from sglang.srt.model_executor.runner_utils import maybe_publish_prefill_war_read_done
 from sglang.srt.speculative.spec_info import SpeculativeAlgorithm
@@ -20,11 +21,12 @@ class _Event:
 
 
 def _model_runner(*, spec_algorithm=SpeculativeAlgorithm.NONE, compliant=True):
+    boundary = (
+        SharedReadBoundary.PRE_REPLAY if compliant else SharedReadBoundary.UNKNOWN
+    )
     return SimpleNamespace(
         spec_algorithm=spec_algorithm,
-        attn_backend=SimpleNamespace(
-            prefill_shared_reads_end_at_metadata_init=compliant
-        ),
+        attn_backend=SimpleNamespace(shared_read_boundary=lambda mode: boundary),
         war_fastpath_read_done_event=None,
     )
 
@@ -60,7 +62,7 @@ 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()),
-            # Backend has not declared metadata-init compliance.
+            # Backend has not declared a pre-replay prefill read boundary.
             (_model_runner(compliant=False), _batch()),
         ):
             maybe_publish_prefill_war_read_done(runner, batch, _DEVICE_MODULE)