From 0f7aaceda5e60a2df56472f87b2b366665215daa Mon Sep 17 00:00:00 2001 From: Liangsheng Yin Date: Sat, 15 Aug 2026 15:02:02 -0700 Subject: [PATCH] [misc] Rename the WAR read-done fastpath to shared-read-done (#34916) --- .../srt/layers/attention/base_attn_backend.py | 2 +- python/sglang/srt/managers/scheduler.py | 6 +- python/sglang/srt/managers/tp_worker.py | 2 +- .../sglang/srt/model_executor/model_runner.py | 13 +- .../runner/decode_cuda_graph_runner.py | 83 +++++---- .../runner/prefill_cuda_graph_runner.py | 6 +- .../model_executor/runner_utils/__init__.py | 5 +- .../{war_event.py => shared_read_event.py} | 12 +- .../srt/speculative/base_spec_worker.py | 2 +- .../eagle_draft_extend_cuda_graph_runner.py | 4 +- .../sglang/srt/speculative/eagle_worker_v2.py | 2 +- python/sglang/srt/speculative/spec_info.py | 4 +- .../sglang/srt/speculative/spec_registry.py | 4 +- ...est_decode_cuda_graph_shared_read_fence.py | 136 ++++++++++++++ .../test_decode_cuda_graph_war_fence.py | 174 ------------------ ...ne.py => test_prefill_shared_read_done.py} | 18 +- 16 files changed, 224 insertions(+), 249 deletions(-) rename python/sglang/srt/model_executor/runner_utils/{war_event.py => shared_read_event.py} (78%) create mode 100644 test/registered/unit/model_executor/runner/test_decode_cuda_graph_shared_read_fence.py delete mode 100644 test/registered/unit/model_executor/runner/test_decode_cuda_graph_war_fence.py rename test/registered/unit/model_executor/runner/{test_prefill_war_read_done.py => test_prefill_shared_read_done.py} (80%) diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index 330385b44..59e7bb508 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -21,7 +21,7 @@ if TYPE_CHECKING: 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 + the shared-read-done record must land at or after this point. IN_REPLAY means at the captured (in-graph) metadata init.""" PRE_REPLAY = auto() diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index f3e73e7a9..12ae8603c 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -1709,9 +1709,9 @@ class Scheduler( # (forceable via SGLANG_FORCE_COARSE_WAR_BARRIER). if not self._war_barrier_enabled: return - runner = self.model_worker.war_fastpath_runner - ev = runner.war_fastpath_read_done_event - runner.war_fastpath_read_done_event = None + runner = self.model_worker.last_shared_read_runner + ev = runner.shared_read_done_event + runner.shared_read_done_event = None if ev is not None and not envs.SGLANG_FORCE_COARSE_WAR_BARRIER.get(): self.schedule_stream.wait_event(ev) else: diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 74efd8ef3..a7a1ef35b 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -82,7 +82,7 @@ class BaseTpWorker(ABC): pass @property - def war_fastpath_runner(self): + def last_shared_read_runner(self): # The runner that runs the step's LAST shared-buffer-reading phase -- # it owns the read-done event the scheduler's WAR barrier waits on. # For a plain worker that's its own runner. diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 72b9c8b3d..075c29371 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -166,7 +166,6 @@ 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, @@ -404,14 +403,10 @@ class ModelRunner: # Init forward stream for overlap schedule self.forward_stream = torch.get_device_module(self.device).Stream() - # WAR fast-path: a decode-graph forward publishes a fresh event here after - # 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) - ) + # 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. + self.shared_read_done_event: Optional[torch.cuda.Event] = None # CPU offload set_offloader( 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 b9c1eae94..5984400f5 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 @@ -93,6 +93,7 @@ from sglang.srt.model_executor.runner_utils.capture_mode import ( from sglang.srt.model_executor.runner_utils.deepep_adapter import ( DeepEPCudaGraphRunnerAdapter, ) +from sglang.srt.model_executor.runner_utils.shared_read_event import make_external_event from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups from sglang.srt.runtime_context import get_flags, get_parallel, get_spec from sglang.srt.speculative.ragged_verify import resolve_ragged_verify_layout @@ -208,7 +209,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): speculative_num_draft_tokens: Optional[int] = None, ): super().__init__(model_runner) - self._war_read_done_node_planted = False + + # In-graph metadata prep: shared buffers -> in-graph private data + self.in_graph_metadata_prep_done: Optional[torch.cuda.Event] = None + # --- core state ------------------------------------------------ self.enable_torch_compile = get_flags().capture.enable_torch_compile self.disable_padding = model_runner.server_args.disable_cuda_graph_padding @@ -424,47 +428,55 @@ 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 _record_in_graph_metadata_prep_done(self): + # Purely a marker at this point in the graph; where the shared reads + # actually end is the attn backend's call. + if not torch.cuda.is_current_stream_capturing(): + # Warmup shares this body. Breakable capture still plants: it opens + # segment 1 on context entry and every segment re-arms the node. + return + if self.in_graph_metadata_prep_done is None: + self.in_graph_metadata_prep_done = make_external_event(self.device_module) + event = self.in_graph_metadata_prep_done + if event is not None: + # Stays None without external-event support, so the boundary + # resolution below never hands out an unrecorded event. + event.record() - def _war_read_done_record(self, attn_backend, forward_mode) -> SharedReadBoundary: - """Where this replay records its WAR read-done event; UNKNOWN records - nothing.""" + def _resolve_shared_read_boundary( + self, attn_backend, forward_mode + ) -> SharedReadBoundary: + """Where this replay records its shared-read-done event: the backend's + declaration, demoted when this runner cannot record at that point. + UNKNOWN records nothing (scheduler keeps the coarse fence).""" if forward_mode.is_target_verify(): - if not self.model_runner.spec_algorithm.is_war_publish_phase(forward_mode): + if not self.model_runner.spec_algorithm.is_last_shared_read_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 + and self.in_graph_metadata_prep_done is None ): - # Non-capturing runs / no external-event support. + # TODO: PRE_REPLAY is EARLIER than the declared boundary; POST_REPLAY + # is the sound demotion for a backend that really reads in-graph. 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.""" + def _publish_read_done(self, in_graph: bool): + """Hand the scheduler's WAR barrier the event marking this phase's + shared-buffer reads as done.""" if in_graph: - self.model_runner.war_fastpath_read_done_event = ( - self.model_runner.war_read_done_event - ) + # Reads end at the in-graph marker: wire it through, don't re-record. + self.model_runner.shared_read_done_event = self.in_graph_metadata_prep_done else: read_done = self.device_module.Event() read_done.record() - self.model_runner.war_fastpath_read_done_event = read_done + self.model_runner.shared_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}) @@ -1066,7 +1078,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() + self._record_in_graph_metadata_prep_done() # No invalidate_loc_cache() here: the unified pool translates its # locs in `init_forward_metadata_out_graph`, so no cache to invalidate. @@ -1304,7 +1316,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): timer_ctx = device_timer_ctx( self.model_runner.device_timer, forward_batch.forward_mode.name.lower() ) - war_record = self._war_read_done_record( + shared_read_boundary = self._resolve_shared_read_boundary( self.attn_backend, forward_batch.forward_mode ) with timer_ctx, self.backend.replay_session(): @@ -1323,13 +1335,16 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): else "" ), ) - if war_record is SharedReadBoundary.PRE_REPLAY: - self._publish_war_read_done(in_graph=False) + if shared_read_boundary is SharedReadBoundary.PRE_REPLAY: + self._publish_read_done(in_graph=False) + output = self.backend.replay(self._replay_graph_key, forward_batch) - if war_record is SharedReadBoundary.POST_REPLAY: - self._publish_war_read_done(in_graph=False) - elif war_record is SharedReadBoundary.IN_REPLAY: - self._publish_war_read_done(in_graph=True) + + if shared_read_boundary is SharedReadBoundary.IN_REPLAY: + self._publish_read_done(in_graph=True) + + if shared_read_boundary is SharedReadBoundary.POST_REPLAY: + self._publish_read_done(in_graph=False) if isinstance(output, LogitsProcessorOutput): if self.is_dllm: diff --git a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py index 908322d2a..167ea8c6a 100644 --- a/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py +++ b/python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py @@ -109,7 +109,9 @@ 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 import ( + maybe_publish_prefill_shared_read_done, +) from sglang.srt.model_executor.runner_utils.buffers import ( PrefillInputBuffers, ) @@ -1758,7 +1760,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner): 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( + maybe_publish_prefill_shared_read_done( self.model_runner, forward_batch, self.device_module ) diff --git a/python/sglang/srt/model_executor/runner_utils/__init__.py b/python/sglang/srt/model_executor/runner_utils/__init__.py index 7332945e0..a8541e9ea 100644 --- a/python/sglang/srt/model_executor/runner_utils/__init__.py +++ b/python/sglang/srt/model_executor/runner_utils/__init__.py @@ -26,7 +26,6 @@ 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 - make_war_read_done_event, - maybe_publish_prefill_war_read_done, +from sglang.srt.model_executor.runner_utils.shared_read_event import ( # noqa: F401 + maybe_publish_prefill_shared_read_done, ) diff --git a/python/sglang/srt/model_executor/runner_utils/war_event.py b/python/sglang/srt/model_executor/runner_utils/shared_read_event.py similarity index 78% rename from python/sglang/srt/model_executor/runner_utils/war_event.py rename to python/sglang/srt/model_executor/runner_utils/shared_read_event.py index f9683f8ae..e117de726 100644 --- a/python/sglang/srt/model_executor/runner_utils/war_event.py +++ b/python/sglang/srt/model_executor/runner_utils/shared_read_event.py @@ -1,4 +1,4 @@ -"""WAR read-done event utilities for CUDA graph runners.""" +"""Shared-read-done event utilities for CUDA graph runners.""" import logging from typing import Optional @@ -13,8 +13,8 @@ from sglang.srt.utils import is_cuda logger = logging.getLogger(__name__) -def make_war_read_done_event(device_module) -> Optional[torch.cuda.Event]: - """Create a persistent external event for CUDA graph capture.""" +def make_external_event(device_module) -> Optional[torch.cuda.Event]: + """Create a persistent external event, e.g., for CUDA graph capture.""" if not is_cuda(): return None try: @@ -23,7 +23,7 @@ def make_war_read_done_event(device_module) -> Optional[torch.cuda.Event]: return None -def maybe_publish_prefill_war_read_done( +def maybe_publish_prefill_shared_read_done( model_runner, forward_batch, device_module ) -> None: """Publish prefill read-done after compliant metadata initialization.""" @@ -42,9 +42,9 @@ def maybe_publish_prefill_war_read_done( if boundary is not SharedReadBoundary.PRE_REPLAY: return logger.info_once( - "Prefill WAR read-done fastpath active (%s)", + "Prefill shared-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 + model_runner.shared_read_done_event = read_done diff --git a/python/sglang/srt/speculative/base_spec_worker.py b/python/sglang/srt/speculative/base_spec_worker.py index 6c2501515..8111c9658 100644 --- a/python/sglang/srt/speculative/base_spec_worker.py +++ b/python/sglang/srt/speculative/base_spec_worker.py @@ -212,7 +212,7 @@ class BaseSpecWorker(ABC): return self.draft_worker.weight_load_time @property - def war_fastpath_runner(self): + def last_shared_read_runner(self): # The runner that runs the step's LAST shared-buffer-reading phase -- # it owns the read-done event the scheduler's WAR barrier waits on. # Default is the target runner; override if the last phase runs diff --git a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py index 02b0cf42c..37b058772 100644 --- a/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py +++ b/python/sglang/srt/speculative/eagle_draft_extend_cuda_graph_runner.py @@ -591,10 +591,10 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner): # Snapshot built -- the forward is done reading the shared pool. Publish # 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). + # is the EAGLE-family last shared-read 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 + self.model_runner.shared_read_done_event = read_done self.raw_bs = raw_bs self.bs = bs diff --git a/python/sglang/srt/speculative/eagle_worker_v2.py b/python/sglang/srt/speculative/eagle_worker_v2.py index fbcabdd11..7d6ddc45a 100644 --- a/python/sglang/srt/speculative/eagle_worker_v2.py +++ b/python/sglang/srt/speculative/eagle_worker_v2.py @@ -1055,7 +1055,7 @@ class EAGLEWorkerV2(BaseSpecWorker): self.plan_stream, self.plan_stream_ctx = get_plan_stream(self.device) @property - def war_fastpath_runner(self): + def last_shared_read_runner(self): # Per the base contract: the step's last shared-buffer-reading phase is # draft_extend, which runs on the draft runner. return self._draft_worker.draft_runner diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index dbe2f566f..e018ee4c0 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/python/sglang/srt/speculative/spec_info.py @@ -127,8 +127,8 @@ class SpeculativeAlgorithm(Enum): def supports_target_verify_for_draft(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 + 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(): diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index 4abbb4103..a736f482a 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/python/sglang/srt/speculative/spec_registry.py @@ -92,8 +92,8 @@ class CustomSpecAlgo: def supports_target_verify_for_draft(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. + 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: 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 new file mode 100644 index 000000000..6a25a3642 --- /dev/null +++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_shared_read_fence.py @@ -0,0 +1,136 @@ +import contextlib +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.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 is_last_shared_read_phase(self, forward_mode) -> bool: + return self._target_verify_war and forward_mode.is_target_verify() + + +def _attn_backend(boundary=SharedReadBoundary.IN_REPLAY): + """Backend stub declaring one fixed read-end boundary for every mode.""" + return SimpleNamespace(shared_read_boundary=lambda _forward_mode: boundary) + + +def _runner(*, target_verify_war: bool = False, has_marker: bool = False): + runner = DecodeCudaGraphRunner.__new__(DecodeCudaGraphRunner) + runner.model_runner = SimpleNamespace( + spec_algorithm=_SpecAlgorithm(target_verify_war), + device_timer=None, + is_draft_worker=False, + shared_read_done_event=None, + ) + runner.in_graph_metadata_prep_done = object() if has_marker else None + return runner + + +def test_unrelated_modes_never_publish(): + # This runner owns the fence for decode / target verify only; every other + # mode stays on the coarse wait, even with a marker available. + assert ( + _runner(has_marker=True)._resolve_shared_read_boundary( + _attn_backend(), ForwardMode.EXTEND + ) + is SharedReadBoundary.UNKNOWN + ) + + +def test_post_replay_declaration_is_not_advanced(): + # A backend that keeps reading shared state across the whole graph declares + # POST_REPLAY. Having an in-graph marker must not pull the fence earlier. + assert ( + _runner(target_verify_war=True, has_marker=True)._resolve_shared_read_boundary( + _attn_backend(SharedReadBoundary.POST_REPLAY), ForwardMode.TARGET_VERIFY + ) + is SharedReadBoundary.POST_REPLAY + ) + + +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_in_graph_marker(): + runner = _runner(has_marker=True) + marker = runner.in_graph_metadata_prep_done + 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.shared_read_done_event is marker + + +def test_execute_falls_back_to_pre_replay_without_marker(): + 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.shared_read_done_event, Event) + + +@pytest.mark.parametrize("supported", [False, True]) +def test_target_verify_requires_war_capability(supported): + runner = _runner(target_verify_war=supported, has_marker=True) + marker = runner.in_graph_metadata_prep_done + runner.attn_backend = _attn_backend() + runner.device_module = SimpleNamespace(Event=lambda: None) + + runner.execute(_execute_harness(runner, [], ForwardMode.TARGET_VERIFY)) + + expected = marker if supported else None + assert runner.model_runner.shared_read_done_event is expected + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_war_fence.py b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_war_fence.py deleted file mode 100644 index bd43843a2..000000000 --- a/test/registered/unit/model_executor/runner/test_decode_cuda_graph_war_fence.py +++ /dev/null @@ -1,174 +0,0 @@ -import contextlib -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.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 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): - 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): - 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_record(): - # Planted node: the graph re-arms it every replay. - assert ( - _runner(planted=True)._war_read_done_record(_attn_backend(), ForwardMode.DECODE) - is SharedReadBoundary.IN_REPLAY - ) - # No planted node: fall back to a pre-replay record. - assert ( - _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_record(_attn_backend(), ForwardMode.EXTEND) - is SharedReadBoundary.UNKNOWN - ) - # The algorithm gate precedes the backend declaration. - assert ( - _runner(planted=True)._war_read_done_record( - _attn_backend(breakable_metadata=True), ForwardMode.TARGET_VERIFY - ) - is SharedReadBoundary.UNKNOWN - ) - # Captured-metadata verify keeps reading throughout the graph, even planted. - assert ( - _runner(target_verify_war=True, planted=True)._war_read_done_record( - _attn_backend(breakable_metadata=True), ForwardMode.TARGET_VERIFY - ) - is SharedReadBoundary.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"])) diff --git a/test/registered/unit/model_executor/runner/test_prefill_war_read_done.py b/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py similarity index 80% rename from test/registered/unit/model_executor/runner/test_prefill_war_read_done.py rename to test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py index 74cb75478..ccac1c1be 100644 --- a/test/registered/unit/model_executor/runner/test_prefill_war_read_done.py +++ b/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py @@ -5,7 +5,9 @@ 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.model_executor.runner_utils import ( + maybe_publish_prefill_shared_read_done, +) from sglang.srt.speculative.spec_info import SpeculativeAlgorithm from sglang.test.ci.ci_register import register_cpu_ci @@ -27,7 +29,7 @@ def _model_runner(*, spec_algorithm=SpeculativeAlgorithm.NONE, compliant=True): return SimpleNamespace( spec_algorithm=spec_algorithm, attn_backend=SimpleNamespace(shared_read_boundary=lambda mode: boundary), - war_fastpath_read_done_event=None, + shared_read_done_event=None, ) @@ -41,16 +43,16 @@ def _batch(mode=ForwardMode.EXTEND): 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 + maybe_publish_prefill_shared_read_done(runner, _batch(), _DEVICE_MODULE) + published = runner.shared_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 + maybe_publish_prefill_shared_read_done(runner, _batch(), _DEVICE_MODULE) + assert runner.shared_read_done_event is None def test_gates_exclude_non_prefill_unsupported_algorithm_and_noncompliant_backend(): @@ -65,8 +67,8 @@ def test_gates_exclude_non_prefill_unsupported_algorithm_and_noncompliant_backen # 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) - assert runner.war_fastpath_read_done_event is None + maybe_publish_prefill_shared_read_done(runner, batch, _DEVICE_MODULE) + assert runner.shared_read_done_event is None if __name__ == "__main__":