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

2026-08-05 13:52:20 -07:00
parent ea65f8ddc9
commit 717a559f02
13 changed files with 509 additions and 24 deletions
@@ -0,0 +1,169 @@
+import contextlib
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+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")
+
+
+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 _attn_backend(*, breakable_metadata=False):
+    return SimpleNamespace(
+        use_captured_forward_metadata_for_breakable_cuda_graph=breakable_metadata,
+    )
+
+
+def _runner(*, target_verify_war: bool = False, planted: bool = False):
+    runner = DecodeCudaGraphRunner.__new__(DecodeCudaGraphRunner)
+    runner.model_runner = SimpleNamespace(
+        spec_algorithm=_SpecAlgorithm(target_verify_war),
+        device_timer=None,
+        is_draft_worker=False,
+        war_read_done_event=None,
+        war_fastpath_read_done_event=None,
+    )
+    runner._war_read_done_node_planted = planted
+    return runner
+
+
+def test_war_read_done_policy():
+    # 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
+    )
+    # No node, snapshot backend: all shared reads finish before launch.
+    assert (
+        _runner()._war_read_done_policy(_attn_backend(), ForwardMode.DECODE)
+        is WarReadDonePolicy.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
+    )
+    # Backend placement cannot opt an unsupported algorithm into publication.
+    assert (
+        _runner(planted=True)._war_read_done_policy(
+            _attn_backend(breakable_metadata=True), ForwardMode.TARGET_VERIFY
+        )
+        is WarReadDonePolicy.NONE
+    )
+    # Captured-metadata verify keeps reading throughout the graph, even planted.
+    assert (
+        _runner(target_verify_war=True, planted=True)._war_read_done_policy(
+            _attn_backend(breakable_metadata=True), ForwardMode.TARGET_VERIFY
+        )
+        is WarReadDonePolicy.POST_REPLAY
+    )
+
+
+def test_publish_war_read_done():
+    runner = _runner()
+    graph_event = object()
+    runner.model_runner.war_read_done_event = graph_event
+    runner._publish_war_read_done(in_graph=True)
+    assert runner.model_runner.war_fastpath_read_done_event is graph_event
+
+    recorded = []
+
+    class Event:
+        def record(self):
+            recorded.append(self)
+
+    runner.device_module = SimpleNamespace(Event=Event)
+    runner._publish_war_read_done(in_graph=False)
+    published = runner.model_runner.war_fastpath_read_done_event
+    assert isinstance(published, Event) and recorded == [published]
+
+
+def _execute_harness(runner, calls, mode=ForwardMode.DECODE):
+    key = ShapeKey(size=1)
+    output = PPProxyTensors({"hidden_states": torch.ones(1, 1)})
+    runner.ragged_verify_mode = False
+    runner.bs = 1
+    runner.load_batch = lambda *_: setattr(runner, "_replay_graph_key", key)
+
+    class Backend:
+        def replay_session(self):
+            return contextlib.nullcontext()
+
+        def replay(self, replay_key, _forward_batch):
+            assert replay_key == key
+            calls.append("replay")
+            return output
+
+    runner.backend = Backend()
+    return SimpleNamespace(forward_mode=mode, batch_size=1)
+
+
+def test_execute_publishes_the_planted_graph_event():
+    runner = _runner(planted=True)
+    graph_event = object()
+    runner.model_runner.war_read_done_event = graph_event
+    runner.attn_backend = _attn_backend()
+    runner.device_module = SimpleNamespace(
+        Event=lambda: (_ for _ in ()).throw(
+            AssertionError("execute must reuse the graph-recorded event")
+        )
+    )
+    calls = []
+    forward_batch = _execute_harness(runner, calls)
+
+    result = runner.execute(forward_batch)
+
+    assert result.tensors["hidden_states"].shape == (1, 1)
+    assert runner.model_runner.war_fastpath_read_done_event is graph_event
+
+
+def test_execute_records_pre_replay_for_snapshot_backends():
+    runner = _runner()
+    runner.attn_backend = _attn_backend()
+    calls = []
+
+    class Event:
+        def record(self):
+            calls.append("record")
+
+    runner.device_module = SimpleNamespace(Event=Event)
+    forward_batch = _execute_harness(runner, calls)
+
+    runner.execute(forward_batch)
+
+    # The eager record lands before the replay so the fence stays truthful.
+    assert calls == ["record", "replay"]
+    assert isinstance(runner.model_runner.war_fastpath_read_done_event, Event)
+
+
+@pytest.mark.parametrize("supported", [False, True])
+def test_target_verify_requires_war_capability(supported):
+    runner = _runner(target_verify_war=supported, planted=True)
+    graph_event = object()
+    runner.model_runner.war_read_done_event = graph_event
+    runner.attn_backend = _attn_backend()
+    runner.device_module = SimpleNamespace(Event=lambda: None)
+
+    runner.execute(_execute_harness(runner, [], ForwardMode.TARGET_VERIFY))
+
+    expected = graph_event if supported else None
+    assert runner.model_runner.war_fastpath_read_done_event is expected
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))
@@ -0,0 +1,71 @@
+from types import SimpleNamespace
+
+import pytest
+
+from sglang.srt.environ import envs
+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
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=1, suite="base-a-test-cpu")
+
+
+class _Event:
+    def __init__(self):
+        self.recorded = False
+
+    def record(self):
+        self.recorded = True
+
+
+def _model_runner(*, spec_algorithm=SpeculativeAlgorithm.NONE, compliant=True):
+    return SimpleNamespace(
+        spec_algorithm=spec_algorithm,
+        attn_backend=SimpleNamespace(
+            prefill_shared_reads_end_at_metadata_init=compliant
+        ),
+        war_fastpath_read_done_event=None,
+    )
+
+
+_DEVICE_MODULE = SimpleNamespace(Event=_Event)
+
+
+def _batch(mode=ForwardMode.EXTEND):
+    return SimpleNamespace(forward_mode=mode)
+
+
+def test_publishes_recorded_event_when_enabled():
+    runner = _model_runner()
+    with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True):
+        maybe_publish_prefill_war_read_done(runner, _batch(), _DEVICE_MODULE)
+    published = runner.war_fastpath_read_done_event
+    assert isinstance(published, _Event) and published.recorded
+
+
+def test_disabled_when_flag_is_false():
+    runner = _model_runner()
+    with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(False):
+        maybe_publish_prefill_war_read_done(runner, _batch(), _DEVICE_MODULE)
+    assert runner.war_fastpath_read_done_event is None
+
+
+def test_gates_exclude_non_prefill_unsupported_algorithm_and_noncompliant_backend():
+    with envs.SGLANG_ENABLE_PREFILL_WAR_READ_DONE.override(True):
+        for runner, batch in (
+            # Verify/mixed/decode publish through the decode graph runner.
+            (_model_runner(), _batch(ForwardMode.TARGET_VERIFY)),
+            (_model_runner(), _batch(ForwardMode.MIXED)),
+            (_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.
+            (_model_runner(compliant=False), _batch()),
+        ):
+            maybe_publish_prefill_war_read_done(runner, batch, _DEVICE_MODULE)
+            assert runner.war_fastpath_read_done_event is None
+
+
+if __name__ == "__main__":
+    raise SystemExit(pytest.main([__file__, "-v"]))