diff --git a/python/sglang/srt/layers/attention/base_attn_backend.py b/python/sglang/srt/layers/attention/base_attn_backend.py index 59e7bb508..614d63216 100644 --- a/python/sglang/srt/layers/attention/base_attn_backend.py +++ b/python/sglang/srt/layers/attention/base_attn_backend.py @@ -1,8 +1,8 @@ from __future__ import annotations from abc import ABC -from enum import Enum, auto -from typing import TYPE_CHECKING, Optional +from enum import Enum +from typing import TYPE_CHECKING, Iterable, Optional import torch @@ -19,15 +19,18 @@ if TYPE_CHECKING: from sglang.srt.speculative.spec_info import SpecInput -class SharedReadBoundary(Enum): - """Where a backend's scheduler-shared reads end, relative to the replay; - the shared-read-done record must land at or after this point. IN_REPLAY - means at the captured (in-graph) metadata init.""" +class SharedReadEnds(Enum): + """Where an attention backend finishes reading the shared data""" - PRE_REPLAY = auto() - IN_REPLAY = auto() - POST_REPLAY = auto() - UNKNOWN = auto() # not audited -> coarse whole-forward fence + PRE_REPLAY = 1 # After the init_forward_metadata_out_graph + IN_REPLAY = 2 # After the init_forward_metadata_in_graph + POST_REPLAY = 3 # Metadata snapshot not implemented + UNKNOWN = 4 # not audited -> coarse whole-forward fence + + @staticmethod + def max_of(items: Iterable[SharedReadEnds]) -> SharedReadEnds: + # Ordered by lateness: the latest end covers every child. + return max(items, key=lambda x: x.value) class AttentionBackend(ABC): @@ -125,16 +128,12 @@ class AttentionBackend(ABC): # object during capture, and refresh its dynamic fields before each replay. use_captured_forward_metadata_for_breakable_cuda_graph: bool = False - def shared_read_boundary(self, forward_mode: ForwardMode) -> SharedReadBoundary: + def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds: """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 + Override only for audited deviations from this conservative default.""" + if fm.is_decode() or fm.is_target_verify(): + return SharedReadEnds.IN_REPLAY + return SharedReadEnds.UNKNOWN # Chunked-prefix FullCG capture has a second model topology and stable # prefix buffers. Backends must opt in explicitly so the runner does not diff --git a/python/sglang/srt/layers/attention/deepseek_v4_backend.py b/python/sglang/srt/layers/attention/deepseek_v4_backend.py index 5830a02a4..e0e4b5377 100644 --- a/python/sglang/srt/layers/attention/deepseek_v4_backend.py +++ b/python/sglang/srt/layers/attention/deepseek_v4_backend.py @@ -41,7 +41,7 @@ from sglang.kernels.ops.speculative.dspark.dspark_attn_metadata import ( from sglang.srt.environ import envs from sglang.srt.layers.attention.base_attn_backend import ( AttentionBackend, - SharedReadBoundary, + SharedReadEnds, ) from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend from sglang.srt.layers.attention.dsv4.compressor_v2 import ( @@ -504,15 +504,15 @@ class DeepseekV4AttnBackend( supports_ragged_verify_graph: bool = True needs_cpu_seq_lens: bool = False - def shared_read_boundary(self, forward_mode: ForwardMode) -> SharedReadBoundary: + def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds: # Breakable-graph verify rereads shared state across segments. # DSPARK verify replays one full (non-breakable) graph that honors the # out-graph/in-graph init contract, so the base IN_REPLAY bound holds. - if forward_mode.is_target_verify(): + if fm.is_target_verify(): if self.model_runner.spec_algorithm.is_dspark(): - return SharedReadBoundary.IN_REPLAY - return SharedReadBoundary.POST_REPLAY - return super().shared_read_boundary(forward_mode) + return SharedReadEnds.IN_REPLAY + return SharedReadEnds.POST_REPLAY + return super().shared_read_ends(fm) def __init__( self, diff --git a/python/sglang/srt/layers/attention/hybrid_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_attn_backend.py index 53b54d949..6c5066659 100644 --- a/python/sglang/srt/layers/attention/hybrid_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_attn_backend.py @@ -4,7 +4,10 @@ from typing import TYPE_CHECKING, Optional import torch -from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.base_attn_backend import ( + AttentionBackend, + SharedReadEnds, +) from sglang.srt.layers.attention.dsa.dsa_indexer_metadata import BaseIndexerMetadata from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode @@ -83,6 +86,9 @@ class HybridAttnBackend(AttentionBackend): else: return self.prefill_backend + def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds: + return self._select_backend(fm).shared_read_ends(fm) + @property def supports_full_cuda_graph_chunked_prefix(self) -> bool: return self.prefill_backend.supports_full_cuda_graph_chunked_prefix diff --git a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py index 08f2eb296..ae8299d00 100644 --- a/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py +++ b/python/sglang/srt/layers/attention/hybrid_linear_attn_backend.py @@ -15,7 +15,10 @@ from sglang.kernels.ops.mamba.mamba_state_scatter_triton import ( track_mamba_states_if_needed, ) from sglang.srt.configs.hybrid_arch import mamba2_config -from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.base_attn_backend import ( + AttentionBackend, + SharedReadEnds, +) from sglang.srt.layers.attention.mamba.mamba import MambaMixer2 from sglang.srt.layers.attention.mamba.mamba2_metadata import ( ForwardMetadata, @@ -1002,6 +1005,11 @@ class HybridLinearAttnBackend(AttentionBackend): forward_batch, in_capture=in_capture ) + def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds: + return SharedReadEnds.max_of( + b.shared_read_ends(fm) for b in self.attn_backend_list + ) + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): for attn_backend in self.attn_backend_list: attn_backend.init_forward_metadata_in_graph(forward_batch) diff --git a/python/sglang/srt/layers/attention/minimax_sparse_backend.py b/python/sglang/srt/layers/attention/minimax_sparse_backend.py index f68ad4b56..def19cbc6 100644 --- a/python/sglang/srt/layers/attention/minimax_sparse_backend.py +++ b/python/sglang/srt/layers/attention/minimax_sparse_backend.py @@ -14,9 +14,12 @@ from sglang.srt.configs.model_config import ( get_minimax_sparse_score_type, ) 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, + SharedReadEnds, +) from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool -from sglang.srt.model_executor.forward_batch_info import ForwardBatch +from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.server_args import m3_fp8_attn_gemm_enabled from sglang.srt.utils import is_npu @@ -1628,6 +1631,11 @@ class MiniMaxHybridAttnBackend(AttentionBackend): self.sparse.init_forward_metadata_out_graph(forward_batch, in_capture) self.dense.init_forward_metadata_out_graph(forward_batch, in_capture) + def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds: + return SharedReadEnds.max_of( + b.shared_read_ends(fm) for b in (self.sparse, self.dense) + ) + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): self.sparse.init_forward_metadata_in_graph(forward_batch) self.dense.init_forward_metadata_in_graph(forward_batch) diff --git a/python/sglang/srt/layers/attention/tbo_backend.py b/python/sglang/srt/layers/attention/tbo_backend.py index 649d40a63..d8713e9f5 100644 --- a/python/sglang/srt/layers/attention/tbo_backend.py +++ b/python/sglang/srt/layers/attention/tbo_backend.py @@ -4,11 +4,14 @@ from types import SimpleNamespace from typing import TYPE_CHECKING, Callable, List, Optional from sglang.srt.batch_overlap import two_batch_overlap -from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.base_attn_backend import ( + AttentionBackend, + SharedReadEnds, +) if TYPE_CHECKING: from sglang.srt.layers.attention.verify_mask import VerifyMask - from sglang.srt.model_executor.forward_batch_info import ForwardBatch + from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode class TboAttnBackend(AttentionBackend): @@ -117,6 +120,11 @@ class TboAttnBackend(AttentionBackend): forward_batch=child_fb_view, in_capture=False ) + def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds: + return SharedReadEnds.max_of( + b.shared_read_ends(fm) for b in (self.primary, *self.children) + ) + def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch): self.primary.init_forward_metadata_in_graph(forward_batch=forward_batch) if not self._children_use_cuda_graph(): diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index 8b86bebf8..34f288b98 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -21,7 +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.base_attn_backend import SharedReadEnds from sglang.srt.layers.attention.flashinfer_backend import ( FlashInferAttnBackend, FlashInferMultiStepDraftBackend, @@ -103,11 +103,11 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): supports_ragged_verify_graph: bool = True - def shared_read_boundary(self, forward_mode: ForwardMode) -> SharedReadBoundary: + def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds: # 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) + if fm == ForwardMode.EXTEND: + return SharedReadEnds.PRE_REPLAY + return super().shared_read_ends(fm) def __init__( self, 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 5984400f5..2af13eca5 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 @@ -44,7 +44,10 @@ 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.base_attn_backend import ( + AttentionBackend, + SharedReadEnds, +) from sglang.srt.layers.attention.dsa.utils import is_dsa_enable_prefill_cp from sglang.srt.layers.dp_attention import ( DpPaddingMode, @@ -439,33 +442,35 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): 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 + # Stays None without external-event support, so the read-end # resolution below never hands out an unrecorded event. event.record() - 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).""" + def _replay_attn_backend(self) -> AttentionBackend: + # Under pdmux each stream replays on its own group member. + if self.enable_pdmux: + return self.model_runner.decode_attn_backend_group[get_current_stream_idx()] + return self.attn_backend + + def _resolve_shared_read_ends(self, attn_backend, forward_mode) -> SharedReadEnds: + """The backend's declaration, demoted when this runner cannot record + there. UNKNOWN records nothing (scheduler keeps the coarse fence).""" if forward_mode.is_target_verify(): if not self.model_runner.spec_algorithm.is_last_shared_read_phase( forward_mode ): - return SharedReadBoundary.UNKNOWN + return SharedReadEnds.UNKNOWN elif not forward_mode.is_decode(): - return SharedReadBoundary.UNKNOWN - boundary = attn_backend.shared_read_boundary(forward_mode) + return SharedReadEnds.UNKNOWN + declared = attn_backend.shared_read_ends(forward_mode) if ( - boundary is SharedReadBoundary.IN_REPLAY + declared is SharedReadEnds.IN_REPLAY and self.in_graph_metadata_prep_done is None ): - # 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 + # TODO: this lands EARLIER than declared; POST_REPLAY is the sound one. + return SharedReadEnds.PRE_REPLAY + return declared def _publish_read_done(self, in_graph: bool): """Hand the scheduler's WAR barrier the event marking this phase's @@ -1271,11 +1276,8 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): and forward_batch.spec_info is not None ): forward_batch.spec_info.custom_mask = buffers.custom_mask - if self.enable_pdmux: - stream_idx = get_current_stream_idx() - attn_backend = self.model_runner.decode_attn_backend_group[stream_idx] - else: - attn_backend = self.attn_backend + + attn_backend = self._replay_attn_backend() fb_view = build_replay_fb_view( forward_batch=forward_batch, buffers=buffers, @@ -1316,8 +1318,8 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): timer_ctx = device_timer_ctx( self.model_runner.device_timer, forward_batch.forward_mode.name.lower() ) - shared_read_boundary = self._resolve_shared_read_boundary( - self.attn_backend, forward_batch.forward_mode + shared_read_ends = self._resolve_shared_read_ends( + self._replay_attn_backend(), forward_batch.forward_mode ) with timer_ctx, self.backend.replay_session(): self.load_batch(forward_batch, pp_proxy_tensors) @@ -1335,15 +1337,15 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner): else "" ), ) - if shared_read_boundary is SharedReadBoundary.PRE_REPLAY: + if shared_read_ends is SharedReadEnds.PRE_REPLAY: self._publish_read_done(in_graph=False) output = self.backend.replay(self._replay_graph_key, forward_batch) - if shared_read_boundary is SharedReadBoundary.IN_REPLAY: + if shared_read_ends is SharedReadEnds.IN_REPLAY: self._publish_read_done(in_graph=True) - if shared_read_boundary is SharedReadBoundary.POST_REPLAY: + if shared_read_ends is SharedReadEnds.POST_REPLAY: self._publish_read_done(in_graph=False) if isinstance(output, LogitsProcessorOutput): diff --git a/python/sglang/srt/model_executor/runner_utils/shared_read_event.py b/python/sglang/srt/model_executor/runner_utils/shared_read_event.py index e117de726..f59e36361 100644 --- a/python/sglang/srt/model_executor/runner_utils/shared_read_event.py +++ b/python/sglang/srt/model_executor/runner_utils/shared_read_event.py @@ -6,7 +6,7 @@ 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.layers.attention.base_attn_backend import SharedReadEnds from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.utils import is_cuda @@ -36,10 +36,8 @@ def maybe_publish_prefill_shared_read_done( if not model_runner.spec_algorithm.is_none(): return # 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: + declared = model_runner.attn_backend.shared_read_ends(forward_batch.forward_mode) + if declared is not SharedReadEnds.PRE_REPLAY: return logger.info_once( "Prefill shared-read-done fastpath active (%s)", 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 index 6a25a3642..5dab628cd 100644 --- 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 @@ -1,135 +1,81 @@ -import contextlib from types import SimpleNamespace +from unittest.mock import create_autospec 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.layers.attention.base_attn_backend import ( + AttentionBackend, + SharedReadEnds, +) +from sglang.srt.model_executor.forward_batch_info import ForwardMode 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() +DECODE = ForwardMode.DECODE +VERIFY = ForwardMode.TARGET_VERIFY +EXTEND = ForwardMode.EXTEND -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): +def _runner(*, owns_verify: 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, + spec_algorithm=SimpleNamespace( + is_last_shared_read_phase=lambda fm: owns_verify and fm.is_target_verify() + ), 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 _backend(declared: SharedReadEnds): + # Spec'd against the real ABC so a rename fails here, not at runtime. + backend = create_autospec(AttentionBackend, instance=True) + backend.shared_read_ends.return_value = declared + return backend -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 - ) +@pytest.mark.parametrize( + "mode, owns_verify, declared, has_marker, expected", + [ + # Only decode / target verify publish; anything else keeps the coarse fence. + (EXTEND, False, SharedReadEnds.IN_REPLAY, True, SharedReadEnds.UNKNOWN), + # Target verify publishes only when it is the step's last reading phase. + (VERIFY, False, SharedReadEnds.IN_REPLAY, True, SharedReadEnds.UNKNOWN), + (VERIFY, True, SharedReadEnds.IN_REPLAY, True, SharedReadEnds.IN_REPLAY), + # A backend that keeps reading through the graph is never advanced. + (VERIFY, True, SharedReadEnds.POST_REPLAY, True, SharedReadEnds.POST_REPLAY), + # Nothing to demote: the declaration is honored as-is. + (DECODE, False, SharedReadEnds.IN_REPLAY, True, SharedReadEnds.IN_REPLAY), + # Nowhere to record in-graph -> fall back to the pre-replay record. + (DECODE, False, SharedReadEnds.IN_REPLAY, False, SharedReadEnds.PRE_REPLAY), + ], +) +def test_resolve_shared_read_ends(mode, owns_verify, declared, has_marker, expected): + runner = _runner(owns_verify=owns_verify, has_marker=has_marker) + assert runner._resolve_shared_read_ends(_backend(declared), mode) is expected -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(): +def test_publish_read_done(): runner = _runner(has_marker=True) - marker = runner.in_graph_metadata_prep_done - runner.attn_backend = _attn_backend() + recorded = [] runner.device_module = SimpleNamespace( - Event=lambda: (_ for _ in ()).throw( - AssertionError("execute must reuse the graph-recorded event") - ) + Event=lambda: SimpleNamespace(record=lambda: recorded.append("record")) ) - 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) + runner._publish_read_done(in_graph=True) + # In-graph: hand over the graph-recorded marker, do not record a new event. marker = runner.in_graph_metadata_prep_done - runner.attn_backend = _attn_backend() - runner.device_module = SimpleNamespace(Event=lambda: None) + assert runner.model_runner.shared_read_done_event is marker + assert recorded == [] - runner.execute(_execute_harness(runner, [], ForwardMode.TARGET_VERIFY)) - - expected = marker if supported else None - assert runner.model_runner.shared_read_done_event is expected + runner._publish_read_done(in_graph=False) + assert recorded == ["record"] + assert runner.model_runner.shared_read_done_event is not marker if __name__ == "__main__": diff --git a/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py b/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py index ccac1c1be..c36e2dd40 100644 --- a/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py +++ b/test/registered/unit/model_executor/runner/test_prefill_shared_read_done.py @@ -1,9 +1,13 @@ from types import SimpleNamespace +from unittest.mock import create_autospec import pytest from sglang.srt.environ import envs -from sglang.srt.layers.attention.base_attn_backend import SharedReadBoundary +from sglang.srt.layers.attention.base_attn_backend import ( + AttentionBackend, + SharedReadEnds, +) from sglang.srt.model_executor.forward_batch_info import ForwardMode from sglang.srt.model_executor.runner_utils import ( maybe_publish_prefill_shared_read_done, @@ -23,12 +27,13 @@ class _Event: def _model_runner(*, spec_algorithm=SpeculativeAlgorithm.NONE, compliant=True): - boundary = ( - SharedReadBoundary.PRE_REPLAY if compliant else SharedReadBoundary.UNKNOWN - ) + declared = SharedReadEnds.PRE_REPLAY if compliant else SharedReadEnds.UNKNOWN + # Spec'd against the real ABC so a rename fails here, not at runtime. + attn_backend = create_autospec(AttentionBackend, instance=True) + attn_backend.shared_read_ends.return_value = declared return SimpleNamespace( spec_algorithm=spec_algorithm, - attn_backend=SimpleNamespace(shared_read_boundary=lambda mode: boundary), + attn_backend=attn_backend, shared_read_done_event=None, ) @@ -64,7 +69,7 @@ def test_gates_exclude_non_prefill_unsupported_algorithm_and_noncompliant_backen (_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 a pre-replay prefill read boundary. + # Backend has not declared a pre-replay prefill read end. (_model_runner(compliant=False), _batch()), ): maybe_publish_prefill_shared_read_done(runner, batch, _DEVICE_MODULE)