From 717a559f02b3ad85ba4bb4623772a1672e9e3e9c Mon Sep 17 00:00:00 2001 From: Jialin Ouyang Date: Wed, 5 Aug 2026 13:52:20 -0700 Subject: [PATCH] [Scheduler] Align WAR fences with CUDA graph metadata reads (#33587) --- python/sglang/srt/environ.py | 2 + .../srt/layers/attention/base_attn_backend.py | 3 + .../layers/attention/trtllm_mha_backend.py | 17 +- .../sglang/srt/model_executor/model_runner.py | 5 + .../runner/decode_cuda_graph_runner.py | 76 +++++--- .../runner/prefill_cuda_graph_runner.py | 6 + .../model_executor/runner_utils/__init__.py | 5 + .../model_executor/runner_utils/war_event.py | 58 ++++++ python/sglang/srt/speculative/spec_info.py | 3 + .../sglang/srt/speculative/spec_registry.py | 3 + .../test_trtllm_mha_graph_metadata.py | 115 ++++++++++++ .../test_decode_cuda_graph_war_fence.py | 169 ++++++++++++++++++ .../runner/test_prefill_war_read_done.py | 71 ++++++++ 13 files changed, 509 insertions(+), 24 deletions(-) create mode 100644 python/sglang/srt/model_executor/runner_utils/war_event.py create mode 100644 test/registered/unit/model_executor/runner/test_decode_cuda_graph_war_fence.py create mode 100644 test/registered/unit/model_executor/runner/test_prefill_war_read_done.py diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 15209060e..907b7876d 100644 --- a/python/sglang/srt/environ.py +++ b/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) diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index 9b7cf373a..bf8bad2e9 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 1583c63a9..1dface94a 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/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( diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 8e957f661..6c6755469 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/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( 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 91e717f21..66bcd6973 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 @@ -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: 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 20e9e83f2..26d23c549 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,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( diff --git a/python/sglang/srt/model_executor/runner_utils/__init__.py b/python/sglang/srt/model_executor/runner_utils/__init__.py index 278c61cae..f0d226f41 100644 --- a/python/sglang/srt/model_executor/runner_utils/__init__.py +++ b/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, +) diff --git a/python/sglang/srt/model_executor/runner_utils/war_event.py b/python/sglang/srt/model_executor/runner_utils/war_event.py new file mode 100644 index 000000000..e8e06aefb --- /dev/null +++ b/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 diff --git a/python/sglang/srt/speculative/spec_info.py b/python/sglang/srt/speculative/spec_info.py index 440bcc0be..7c1f8b98c 100644 --- a/python/sglang/srt/speculative/spec_info.py +++ b/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 diff --git a/python/sglang/srt/speculative/spec_registry.py b/python/sglang/srt/speculative/spec_registry.py index 4ab16a2a3..d2ac890a0 100644 --- a/python/sglang/srt/speculative/spec_registry.py +++ b/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 diff --git a/test/registered/attention/test_trtllm_mha_graph_metadata.py b/test/registered/attention/test_trtllm_mha_graph_metadata.py index f9776949c..83a08ef09 100644 --- a/test/registered/attention/test_trtllm_mha_graph_metadata.py +++ b/test/registered/attention/test_trtllm_mha_graph_metadata.py @@ -207,6 +207,121 @@ def test_metadata_update_records_inside_cuda_graph(): ) +def test_graph_read_done_event_fences_slot_mutation(): + if not torch.cuda.is_available(): + pytest.skip("CUDA required") + + backend = _make_backend_for_hook_test() + backend.device = torch.device(DEVICE) + backend.page_size = 2 + backend.max_num_pages = 2 + backend.use_sliding_window_kv_pool = True + backend._swa_kv_pool = object() + backend.req_to_token = torch.tensor( + [[0, 1, 2, 3], [8, 9, 10, 11]], dtype=torch.int32, device=DEVICE + ) + backend._swa_full_to_swa_mapping = ( + torch.arange(32, dtype=torch.int64, device=DEVICE) * 2 + ) + backend.init_cuda_graph_state(max_bs=1, max_num_tokens=1) + forward_batch = SimpleNamespace( + batch_size=1, + req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE), + seq_lens=torch.tensor([4], dtype=torch.int32, device=DEVICE), + forward_mode=ForwardMode.DECODE, + spec_info=None, + positions=torch.tensor([0], dtype=torch.int64, device=DEVICE), + out_cache_loc=torch.tensor([3], dtype=torch.int64, device=DEVICE), + ) + + backend.init_forward_metadata_out_graph(forward_batch, in_capture=True) + backend.init_forward_metadata_in_graph(forward_batch) + torch.cuda.synchronize() + + read_done = torch.cuda.Event(external=True) + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + backend.init_forward_metadata_in_graph(forward_batch) + read_done.record() + + fence_stream = torch.cuda.Stream() + mutation_done = torch.cuda.Event() + graph.replay() + with torch.cuda.stream(fence_stream): + fence_stream.wait_event(read_done) + backend.req_to_token.copy_( + torch.tensor( + [[8, 9, 10, 11], [0, 1, 2, 3]], + dtype=torch.int32, + device=DEVICE, + ) + ) + backend._swa_full_to_swa_mapping.add_(64) + mutation_done.record() + torch.cuda.current_stream().wait_event(mutation_done) + torch.cuda.synchronize() + + torch.testing.assert_close( + backend.forward_metadata.page_table, + torch.tensor([[0, 1]], dtype=torch.int32, device=DEVICE), + ) + torch.testing.assert_close( + backend.forward_metadata.swa_page_table, + torch.tensor([[0, 2]], dtype=torch.int32, device=DEVICE), + ) + torch.testing.assert_close( + backend.forward_metadata.swa_out_cache_loc, + torch.tensor([6], dtype=torch.int64, device=DEVICE), + ) + + graph.replay() + with torch.cuda.stream(fence_stream): + fence_stream.wait_event(read_done) + backend.req_to_token.copy_( + torch.tensor( + [[0, 1, 2, 3], [8, 9, 10, 11]], + dtype=torch.int32, + device=DEVICE, + ) + ) + backend._swa_full_to_swa_mapping.sub_(64) + mutation_done.record() + torch.cuda.current_stream().wait_event(mutation_done) + torch.cuda.synchronize() + + torch.testing.assert_close( + backend.forward_metadata.page_table, + torch.tensor([[4, 5]], dtype=torch.int32, device=DEVICE), + ) + torch.testing.assert_close( + backend.forward_metadata.swa_page_table, + torch.tensor([[40, 42]], dtype=torch.int32, device=DEVICE), + ) + torch.testing.assert_close( + backend.forward_metadata.swa_out_cache_loc, + torch.tensor([70], dtype=torch.int64, device=DEVICE), + ) + + +def test_swa_cache_write_uses_metadata_slot_snapshot(): + snapshot = torch.tensor([6], dtype=torch.int64) + + def translate_live_mapping(_): + raise AssertionError("cache writes must not read the live SWA mapping") + + backend = TRTLLMHAAttnBackend.__new__(TRTLLMHAAttnBackend) + backend._swa_kv_pool = SimpleNamespace( + layers_mapping={1: (0, True)}, + translate_loc_from_full_to_swa=translate_live_mapping, + ) + backend.forward_metadata = SimpleNamespace(swa_out_cache_loc=snapshot) + forward_batch = SimpleNamespace(out_cache_loc=torch.tensor([3], dtype=torch.int64)) + + cache_loc = backend._get_layer_cache_loc(SimpleNamespace(layer_id=1), forward_batch) + + torch.testing.assert_close(cache_loc, snapshot, rtol=0, atol=0) + + def _build_inputs(bs, pool_size, max_num_pages, max_seq_pages, seq_max, seed): """Build random pool / indices / seq_lens consistent with backend buffers.""" g = torch.Generator(device="cpu").manual_seed(seed) 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 new file mode 100644 index 000000000..ee3cfda14 --- /dev/null +++ b/test/registered/unit/model_executor/runner/test_decode_cuda_graph_war_fence.py @@ -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"])) diff --git a/test/registered/unit/model_executor/runner/test_prefill_war_read_done.py b/test/registered/unit/model_executor/runner/test_prefill_war_read_done.py new file mode 100644 index 000000000..7972edaef --- /dev/null +++ b/test/registered/unit/model_executor/runner/test_prefill_war_read_done.py @@ -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"]))