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

python/sglang/srt/layers/attention/base_attn_backend.py

@@ -1,6 +1,7 @@
 from __future__ import annotations
 
 from abc import ABC
+from enum import Enum, auto
 from typing import TYPE_CHECKING, Optional
 
 import torch
@@ -14,10 +15,21 @@ if TYPE_CHECKING:
     )
     from sglang.srt.layers.attention.verify_mask import VerifyMask
     from sglang.srt.layers.radix_attention import RadixAttention
-    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
     from sglang.srt.speculative.spec_info import SpecInput
 
 
+class SharedReadBoundary(Enum):
+    """Where a backend's scheduler-shared reads end, relative to the replay;
+    the WAR read-done record must land at or after this point. IN_REPLAY
+    means at the captured (in-graph) metadata init."""
+
+    PRE_REPLAY = auto()
+    IN_REPLAY = auto()
+    POST_REPLAY = auto()
+    UNKNOWN = auto()  # not audited -> coarse whole-forward fence
+
+
 class AttentionBackend(ABC):
     """The base class of attention backends.
 
@@ -113,8 +125,16 @@ class AttentionBackend(ABC):
     # object during capture, and refresh its dynamic fields before each replay.
     use_captured_forward_metadata_for_breakable_cuda_graph: bool = False
 
-    # Whether prefill metadata initialization finishes all scheduler-shared reads.
-    prefill_shared_reads_end_at_metadata_init: bool = False
+    def shared_read_boundary(self, forward_mode: ForwardMode) -> SharedReadBoundary:
+        """Declare where this backend's scheduler-shared reads end per mode.
+
+        Decode/verify default to IN_REPLAY: the out-graph/in-graph init
+        contract above makes it a safe upper bound for any backend honoring
+        the contract. Override for audited deviations.
+        """
+        if forward_mode.is_decode() or forward_mode.is_target_verify():
+            return SharedReadBoundary.IN_REPLAY
+        return SharedReadBoundary.UNKNOWN
 
     # Chunked-prefix FullCG capture has a second model topology and stable
     # prefix buffers. Backends must opt in explicitly so the runner does not

python/sglang/srt/layers/attention/deepseek_v4_backend.py

@@ -39,7 +39,10 @@ from sglang.kernels.ops.speculative.dspark.dspark_attn_metadata import (
     ComputeDsparkWindowGather,
 )
 from sglang.srt.environ import envs
-from sglang.srt.layers.attention.base_attn_backend import AttentionBackend
+from sglang.srt.layers.attention.base_attn_backend import (
+    AttentionBackend,
+    SharedReadBoundary,
+)
 from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend
 from sglang.srt.layers.attention.dsv4.compressor_v2 import (
     CompressorBackendMixin,
@@ -501,6 +504,12 @@ class DeepseekV4AttnBackend(
     supports_ragged_verify_graph: bool = True
     needs_cpu_seq_lens: bool = False
 
+    def shared_read_boundary(self, forward_mode: ForwardMode) -> SharedReadBoundary:
+        # Breakable-graph verify rereads shared state across segments.
+        if forward_mode.is_target_verify():
+            return SharedReadBoundary.POST_REPLAY
+        return super().shared_read_boundary(forward_mode)
+
     def __init__(
         self,
         model_runner: ModelRunner,

python/sglang/srt/layers/attention/trtllm_mha_backend.py

@@ -21,6 +21,7 @@ from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
     build_trtllm_mha_page_table,
 )
 from sglang.srt.environ import envs
+from sglang.srt.layers.attention.base_attn_backend import SharedReadBoundary
 from sglang.srt.layers.attention.flashinfer_backend import (
     FlashInferAttnBackend,
     FlashInferMultiStepDraftBackend,
@@ -102,8 +103,11 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
 
     supports_ragged_verify_graph: bool = True
 
-    # Prefill metadata initialization snapshots all scheduler-shared inputs.
-    prefill_shared_reads_end_at_metadata_init: bool = True
+    def shared_read_boundary(self, forward_mode: ForwardMode) -> SharedReadBoundary:
+        # Prefill metadata init snapshots all scheduler-shared inputs pre-replay.
+        if forward_mode == ForwardMode.EXTEND:
+            return SharedReadBoundary.PRE_REPLAY
+        return super().shared_read_boundary(forward_mode)
 
     def __init__(
         self,

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -44,6 +44,7 @@ from sglang.srt.distributed.parallel_state import (
 )
 from sglang.srt.dllm.config import DllmConfig
 from sglang.srt.environ import envs
+from sglang.srt.layers.attention.base_attn_backend import SharedReadBoundary
 from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp
 from sglang.srt.layers.dp_attention import (
     DpPaddingMode,
@@ -82,7 +83,6 @@ from sglang.srt.model_executor.runner_backend.utils import resolve_decode_backen
 from sglang.srt.model_executor.runner_backend_utils import (
     CUDA_GRAPH_CAPTURE_FAILED_MSG,
 )
-from sglang.srt.model_executor.runner_utils import WarReadDonePolicy
 from sglang.srt.model_executor.runner_utils.buffers import (
     DecodeInputBuffers,
 )
@@ -438,23 +438,22 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             self.model_runner.war_read_done_event.record()
             self._war_read_done_node_planted = True
 
-    def _war_read_done_policy(self, attn_backend, forward_mode) -> WarReadDonePolicy:
-        """Whether and where this replay records its WAR read-done event."""
+    def _war_read_done_record(self, attn_backend, forward_mode) -> SharedReadBoundary:
+        """Where this replay records its WAR read-done event; UNKNOWN records
+        nothing."""
         if forward_mode.is_target_verify():
-            if (
-                not self.model_runner.spec_algorithm.supports_target_verify_war_read_done()
-            ):
-                return WarReadDonePolicy.NONE
-            if attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph:
-                return WarReadDonePolicy.POST_REPLAY
-            if self._war_read_done_node_planted:
-                return WarReadDonePolicy.IN_GRAPH
-            return WarReadDonePolicy.PRE_REPLAY
-        elif forward_mode.is_decode():
-            if self._war_read_done_node_planted:
-                return WarReadDonePolicy.IN_GRAPH
-            return WarReadDonePolicy.PRE_REPLAY
-        return WarReadDonePolicy.NONE
+            if not self.model_runner.spec_algorithm.is_war_publish_phase(forward_mode):
+                return SharedReadBoundary.UNKNOWN
+        elif not forward_mode.is_decode():
+            return SharedReadBoundary.UNKNOWN
+        boundary = attn_backend.shared_read_boundary(forward_mode)
+        if (
+            boundary is SharedReadBoundary.IN_REPLAY
+            and not self._war_read_done_node_planted
+        ):
+            # Non-capturing runs / no external-event support.
+            return SharedReadBoundary.PRE_REPLAY
+        return boundary
 
     def _publish_war_read_done(self, in_graph: bool):
         """Publish the read-done event the scheduler's WAR barrier waits on."""
@@ -1305,7 +1304,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         timer_ctx = device_timer_ctx(
             self.model_runner.device_timer, forward_batch.forward_mode.name.lower()
         )
-        war_policy = self._war_read_done_policy(
+        war_record = self._war_read_done_record(
             self.attn_backend, forward_batch.forward_mode
         )
         with timer_ctx, self.backend.replay_session():
@@ -1324,12 +1323,12 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                         else ""
                     ),
                 )
-            if war_policy is WarReadDonePolicy.PRE_REPLAY:
+            if war_record is SharedReadBoundary.PRE_REPLAY:
                 self._publish_war_read_done(in_graph=False)
             output = self.backend.replay(self._replay_graph_key, forward_batch)
-            if war_policy is WarReadDonePolicy.POST_REPLAY:
+            if war_record is SharedReadBoundary.POST_REPLAY:
                 self._publish_war_read_done(in_graph=False)
-            elif war_policy is WarReadDonePolicy.IN_GRAPH:
+            elif war_record is SharedReadBoundary.IN_REPLAY:
                 self._publish_war_read_done(in_graph=True)
 
         if isinstance(output, LogitsProcessorOutput):

python/sglang/srt/model_executor/runner_utils/__init__.py

@@ -27,7 +27,6 @@ from sglang.srt.model_executor.runner_utils.pool import (  # noqa: F401
     set_global_graph_memory_pool,
 )
 from sglang.srt.model_executor.runner_utils.war_event import (  # noqa: F401
-    WarReadDonePolicy,
     make_war_read_done_event,
     maybe_publish_prefill_war_read_done,
 )

python/sglang/srt/model_executor/runner_utils/war_event.py

@@ -1,30 +1,18 @@
 """WAR read-done event utilities for CUDA graph runners."""
 
 import logging
-from enum import Enum, auto
 from typing import Optional
 
 import torch
 
 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.utils import is_cuda
 
 logger = logging.getLogger(__name__)
 
 
-# Whether and where the WAR read-done record lands for a replay.
-class WarReadDonePolicy(Enum):
-    # This forward mode or algorithm does not publish a read-done event.
-    NONE = auto()
-    # Snapshot backends finish all shared reads before launch.
-    PRE_REPLAY = auto()
-    # Captured metadata initialization finishes all shared reads.
-    IN_GRAPH = auto()
-    # Captured-metadata replays keep reading shared buffers throughout the graph.
-    POST_REPLAY = auto()
-
-
 def make_war_read_done_event(device_module) -> Optional[torch.cuda.Event]:
     """Create a persistent external event for CUDA graph capture."""
     if not is_cuda():
@@ -47,7 +35,11 @@ def maybe_publish_prefill_war_read_done(
     # WAR boundaries are validated.
     if not model_runner.spec_algorithm.is_none():
         return
-    if not model_runner.attn_backend.prefill_shared_reads_end_at_metadata_init:
+    # The record lands right after replay prep, so PRE_REPLAY only.
+    boundary = model_runner.attn_backend.shared_read_boundary(
+        forward_batch.forward_mode
+    )
+    if boundary is not SharedReadBoundary.PRE_REPLAY:
         return
     logger.info_once(
         "Prefill WAR read-done fastpath active (%s)",

python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py

@@ -590,7 +590,8 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
         self.draft_extend_attn_backend.init_forward_metadata_out_graph(fb_view)
 
         # Snapshot built -- the forward is done reading the shared pool. Publish
-        # a read-done event the scheduler's WAR barrier waits on.
+        # a read-done event the scheduler's WAR barrier waits on (draft extend
+        # is the EAGLE-family war-publish phase; last write wins the mailbox).
         read_done = self.device_module.Event()
         read_done.record()
         self.model_runner.war_fastpath_read_done_event = read_done

python/sglang/srt/speculative/spec_info.py

@@ -127,8 +127,11 @@ class SpeculativeAlgorithm(Enum):
     def supports_target_verify_for_draft(self) -> bool:
         return self.is_dflash_family()
 
-    def supports_target_verify_war_read_done(self) -> bool:
-        return self.is_dflash_family()
+    def is_war_publish_phase(self, forward_mode) -> bool:
+        # The step's last shared-buffer-reading phase owns the WAR read-done publish.
+        if self.is_dflash_family():
+            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

python/sglang/srt/speculative/spec_registry.py

@@ -92,8 +92,9 @@ class CustomSpecAlgo:
     def supports_target_verify_for_draft(self) -> bool:
         return False
 
-    def supports_target_verify_war_read_done(self) -> bool:
-        return False
+    def is_war_publish_phase(self, forward_mode) -> bool:
+        # The step's last shared-buffer-reading phase owns the WAR read-done publish.
+        return forward_mode.is_draft_extend_v2()
 
     def supports_ragged_verify(self) -> bool:
         return False