[Scheduler] Align WAR fences with CUDA graph metadata reads (#33587)

python/sglang/srt/environ.py

@@ -433,6 +433,8 @@ class Envs:
     # Force the WAR barrier to wait for the whole forward instead of the
     # read-done fastpath event.
     SGLANG_FORCE_COARSE_WAR_BARRIER = EnvBool(False)
+    # Enable prefill read-done publication after compliant metadata initialization.
+    SGLANG_ENABLE_PREFILL_WAR_READ_DONE = EnvBool(False)
     # PP: skip output send/recv when the entire batch consists of non-final chunked prefill requests,
     # since process_batch_result_prefill discards next_token_ids for those anyway.
     SGLANG_PP_SKIP_PURE_CHUNKED_OUTPUT_COMM = EnvBool(False)

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

@@ -108,6 +108,9 @@ 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
+
     # Chunked-prefix FullCG capture has a second model topology and stable
     # prefix buffers. Backends must opt in explicitly so the runner does not
     # assume that generic ForwardBatch metadata is sufficient for every

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

@@ -102,6 +102,9 @@ 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 __init__(
         self,
         model_runner: ModelRunner,
@@ -340,9 +343,19 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
         if self._swa_kv_pool is not None:
             _, is_swa = self._swa_kv_pool.layers_mapping[layer.layer_id]
             if is_swa:
-                return self._swa_kv_pool.translate_loc_from_full_to_swa(
-                    forward_batch.out_cache_loc
+                swa_loc = self.forward_metadata.swa_out_cache_loc
+                assert (
+                    swa_loc is not None
+                    and swa_loc.shape[0] >= forward_batch.out_cache_loc.shape[0]
+                ), (
+                    "SWA write locs missing or too short: init_forward_metadata "
+                    "must translate out_cache_loc once; a per-layer gather of "
+                    "the live mapping would race the scheduler after the WAR "
+                    "fence releases"
                 )
+                # Piecewise prefill narrows out_cache_loc per attention call;
+                # the snapshot keeps the padded batch length.
+                return swa_loc[: forward_batch.out_cache_loc.shape[0]]
         return forward_batch.out_cache_loc
 
     def _bind_swa_page_table(

python/sglang/srt/model_executor/model_runner.py

@@ -166,6 +166,7 @@ from sglang.srt.model_executor.runner import (
     EagerRunner,
     get_batch_sizes_to_capture,
 )
+from sglang.srt.model_executor.runner_utils import make_war_read_done_event
 from sglang.srt.platforms import current_platform
 from sglang.srt.runtime_context import (
     get_context,
@@ -374,6 +375,10 @@ class ModelRunner:
         # load_batch; the scheduler's WAR barrier waits on it (then clears it)
         # instead of the whole-forward wait_stream. None -> whole-forward fallback.
         self.war_fastpath_read_done_event: Optional[torch.cuda.Event] = None
+        # Graph runners record this persistent event after shared-state reads.
+        self.war_read_done_event = make_war_read_done_event(
+            torch.get_device_module(self.device)
+        )
 
         # CPU offload
         set_offloader(

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

@@ -81,6 +81,7 @@ 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,
 )
@@ -203,6 +204,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         speculative_num_draft_tokens: Optional[int] = None,
     ):
         super().__init__(model_runner)
+        self._war_read_done_node_planted = False
         # --- core state ------------------------------------------------
         self.enable_torch_compile = get_flags().capture.enable_torch_compile
         self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
@@ -418,6 +420,49 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 f"Capture cuda graph failed: {e}\n" f"{CUDA_GRAPH_CAPTURE_FAILED_MSG}"
             )
 
+    def _plant_war_read_done_node(self):
+        """Record the read-done event as a graph node right after the in-graph
+        metadata hook. Safe for both full and breakable capture: capture is
+        active here (breakable opens segment 1 on context entry) and every
+        segment replays, re-arming the node. Non-capturing runs (warmup,
+        debug-eager, no external-event support) leave the flag unset and stay
+        on the fallback paths."""
+        if (
+            self.model_runner.war_read_done_event is not None
+            and torch.cuda.is_current_stream_capturing()
+        ):
+            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."""
+        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
+
+    def _publish_war_read_done(self, in_graph: bool):
+        """Publish the read-done event the scheduler's WAR barrier waits on."""
+        if in_graph:
+            self.model_runner.war_fastpath_read_done_event = (
+                self.model_runner.war_read_done_event
+            )
+        else:
+            read_done = self.device_module.Event()
+            read_done.record()
+            self.model_runner.war_fastpath_read_done_event = read_done
+
     def _build_ragged_verify_token_buckets(self) -> list[int]:
         buckets = sorted({bs * self.captured_req_width for bs in self.capture_bs})
         assert buckets and buckets[0] > 0, f"{buckets=}"
@@ -958,6 +1003,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                 # run eagerly in `init_forward_metadata_out_graph` (replay-prep), so
                 # the captured graph reads already-physical locs. Base no-op for triton.
                 attn_backend.init_forward_metadata_in_graph(forward_batch)
+                self._plant_war_read_done_node()
 
                 # No invalidate_loc_cache() here: the unified pool translates its
                 # locs in `init_forward_metadata_out_graph`, so no cache to invalidate.
@@ -1195,20 +1241,8 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
         timer_ctx = device_timer_ctx(
             self.model_runner.device_timer, forward_batch.forward_mode.name.lower()
         )
-        # Publish a read-done event for the WAR barrier: a cuda-graph forward
-        # finishes its shared req_to_token / SWA reads at this pre-replay
-        # snapshot, so plain DECODE and block-draft TARGET_VERIFY qualify.
-        publish_read_done = forward_batch.forward_mode.is_decode() or (
-            forward_batch.forward_mode.is_target_verify()
-            and self.model_runner.spec_algorithm.is_dflash_family()
-        )
-        # Exception: breakable-graph verify replays (captured forward metadata)
-        # re-read req_to_token *during* replay, so the pre-replay snapshot is
-        # too early -- record the event after replay instead.
-        read_done_post_replay = (
-            publish_read_done
-            and forward_batch.forward_mode.is_target_verify()
-            and self.attn_backend.use_captured_forward_metadata_for_breakable_cuda_graph
+        war_policy = self._war_read_done_policy(
+            self.attn_backend, forward_batch.forward_mode
         )
         with timer_ctx, self.backend.replay_session():
             self.load_batch(forward_batch, pp_proxy_tensors)
@@ -1226,15 +1260,13 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
                         else ""
                     ),
                 )
-            if publish_read_done and not read_done_post_replay:
-                read_done = self.device_module.Event()
-                read_done.record()
-                self.model_runner.war_fastpath_read_done_event = read_done
+            if war_policy is WarReadDonePolicy.PRE_REPLAY:
+                self._publish_war_read_done(in_graph=False)
             output = self.backend.replay(self._replay_graph_key, forward_batch)
-            if read_done_post_replay:
-                read_done = self.device_module.Event()
-                read_done.record()
-                self.model_runner.war_fastpath_read_done_event = read_done
+            if war_policy is WarReadDonePolicy.POST_REPLAY:
+                self._publish_war_read_done(in_graph=False)
+            elif war_policy is WarReadDonePolicy.IN_GRAPH:
+                self._publish_war_read_done(in_graph=True)
 
         if isinstance(output, LogitsProcessorOutput):
             if self.is_dllm:

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

@@ -109,6 +109,7 @@ from sglang.srt.model_executor.runner_backend_utils.tc_piecewise_cuda_graph impo
     TCPCG_FAILURE_HINT,
     set_tc_piecewise_forward_context,
 )
+from sglang.srt.model_executor.runner_utils import maybe_publish_prefill_war_read_done
 from sglang.srt.model_executor.runner_utils.buffers import (
     PrefillInputBuffers,
 )
@@ -1737,6 +1738,11 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
             # The only variants this runner records are chunked-prefix ones.
             if shape_key.variant_label is not None:
                 self._prepare_chunked_prefix_replay(shape_key, forward_batch)
+            # Replay prep, including the optional chunked-prefix gather above,
+            # has finished every scheduler-shared read.
+            maybe_publish_prefill_war_read_done(
+                self.model_runner, forward_batch, self.device_module
+            )
 
             if self.enable_cp_v2_bcg_capture:
                 output = execute_prefill_cp_bcg(

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

@@ -26,3 +26,8 @@ from sglang.srt.model_executor.runner_utils.pool import (  # noqa: F401
     get_global_graph_memory_pool,
     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

@@ -0,0 +1,58 @@
+"""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.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():
+        return None
+    try:
+        return device_module.Event(external=True)
+    except TypeError:
+        return None
+
+
+def maybe_publish_prefill_war_read_done(
+    model_runner, forward_batch, device_module
+) -> None:
+    """Publish prefill read-done after compliant metadata initialization."""
+    if not envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.get():
+        return
+    if forward_batch.forward_mode != ForwardMode.EXTEND:
+        return
+    # TODO(Jialin): Relax this gate for speculative decoding after its prefill
+    # 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:
+        return
+    logger.info_once(
+        "Prefill WAR read-done fastpath active (%s)",
+        type(model_runner.attn_backend).__name__,
+    )
+    read_done = device_module.Event()
+    read_done.record()
+    model_runner.war_fastpath_read_done_event = read_done

python/sglang/srt/speculative/spec_info.py

@@ -127,6 +127,9 @@ 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 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

python/sglang/srt/speculative/spec_registry.py

@@ -92,6 +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 supports_ragged_verify(self) -> bool:
         return False