[misc] Rename shared-read boundary to shared-read ends and fix wrapper delegation (#34982)

This commit is contained in:
Liangsheng Yin
2026-08-16 14:36:31 -07:00
committed by GitHub
parent ace7314173
commit bae353ba55
11 changed files with 154 additions and 174 deletions
@@ -1,8 +1,8 @@
from __future__ import annotations from __future__ import annotations
from abc import ABC from abc import ABC
from enum import Enum, auto from enum import Enum
from typing import TYPE_CHECKING, Optional from typing import TYPE_CHECKING, Iterable, Optional
import torch import torch
@@ -19,15 +19,18 @@ if TYPE_CHECKING:
from sglang.srt.speculative.spec_info import SpecInput from sglang.srt.speculative.spec_info import SpecInput
class SharedReadBoundary(Enum): class SharedReadEnds(Enum):
"""Where a backend's scheduler-shared reads end, relative to the replay; """Where an attention backend finishes reading the shared data"""
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() PRE_REPLAY = 1 # After the init_forward_metadata_out_graph
IN_REPLAY = auto() IN_REPLAY = 2 # After the init_forward_metadata_in_graph
POST_REPLAY = auto() POST_REPLAY = 3 # Metadata snapshot not implemented
UNKNOWN = auto() # not audited -> coarse whole-forward fence 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): class AttentionBackend(ABC):
@@ -125,16 +128,12 @@ class AttentionBackend(ABC):
# object during capture, and refresh its dynamic fields before each replay. # object during capture, and refresh its dynamic fields before each replay.
use_captured_forward_metadata_for_breakable_cuda_graph: bool = False 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. """Declare where this backend's scheduler-shared reads end per mode.
Override only for audited deviations from this conservative default."""
Decode/verify default to IN_REPLAY: the out-graph/in-graph init if fm.is_decode() or fm.is_target_verify():
contract above makes it a safe upper bound for any backend honoring return SharedReadEnds.IN_REPLAY
the contract. Override for audited deviations. return SharedReadEnds.UNKNOWN
"""
if forward_mode.is_decode() or forward_mode.is_target_verify():
return SharedReadBoundary.IN_REPLAY
return SharedReadBoundary.UNKNOWN
# Chunked-prefix FullCG capture has a second model topology and stable # Chunked-prefix FullCG capture has a second model topology and stable
# prefix buffers. Backends must opt in explicitly so the runner does not # prefix buffers. Backends must opt in explicitly so the runner does not
@@ -41,7 +41,7 @@ from sglang.kernels.ops.speculative.dspark.dspark_attn_metadata import (
from sglang.srt.environ import envs from sglang.srt.environ import envs
from sglang.srt.layers.attention.base_attn_backend import ( from sglang.srt.layers.attention.base_attn_backend import (
AttentionBackend, AttentionBackend,
SharedReadBoundary, SharedReadEnds,
) )
from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend from sglang.srt.layers.attention.dsa.dsa_topk_backend import DSATopKBackend
from sglang.srt.layers.attention.dsv4.compressor_v2 import ( from sglang.srt.layers.attention.dsv4.compressor_v2 import (
@@ -504,15 +504,15 @@ class DeepseekV4AttnBackend(
supports_ragged_verify_graph: bool = True supports_ragged_verify_graph: bool = True
needs_cpu_seq_lens: bool = False 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. # Breakable-graph verify rereads shared state across segments.
# DSPARK verify replays one full (non-breakable) graph that honors the # DSPARK verify replays one full (non-breakable) graph that honors the
# out-graph/in-graph init contract, so the base IN_REPLAY bound holds. # 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(): if self.model_runner.spec_algorithm.is_dspark():
return SharedReadBoundary.IN_REPLAY return SharedReadEnds.IN_REPLAY
return SharedReadBoundary.POST_REPLAY return SharedReadEnds.POST_REPLAY
return super().shared_read_boundary(forward_mode) return super().shared_read_ends(fm)
def __init__( def __init__(
self, self,
@@ -4,7 +4,10 @@ from typing import TYPE_CHECKING, Optional
import torch 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.attention.dsa.dsa_indexer_metadata import BaseIndexerMetadata
from sglang.srt.layers.radix_attention import RadixAttention from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
@@ -83,6 +86,9 @@ class HybridAttnBackend(AttentionBackend):
else: else:
return self.prefill_backend return self.prefill_backend
def shared_read_ends(self, fm: ForwardMode) -> SharedReadEnds:
return self._select_backend(fm).shared_read_ends(fm)
@property @property
def supports_full_cuda_graph_chunked_prefix(self) -> bool: def supports_full_cuda_graph_chunked_prefix(self) -> bool:
return self.prefill_backend.supports_full_cuda_graph_chunked_prefix return self.prefill_backend.supports_full_cuda_graph_chunked_prefix
@@ -15,7 +15,10 @@ from sglang.kernels.ops.mamba.mamba_state_scatter_triton import (
track_mamba_states_if_needed, track_mamba_states_if_needed,
) )
from sglang.srt.configs.hybrid_arch import mamba2_config 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.mamba import MambaMixer2
from sglang.srt.layers.attention.mamba.mamba2_metadata import ( from sglang.srt.layers.attention.mamba.mamba2_metadata import (
ForwardMetadata, ForwardMetadata,
@@ -1002,6 +1005,11 @@ class HybridLinearAttnBackend(AttentionBackend):
forward_batch, in_capture=in_capture 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): def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch):
for attn_backend in self.attn_backend_list: for attn_backend in self.attn_backend_list:
attn_backend.init_forward_metadata_in_graph(forward_batch) attn_backend.init_forward_metadata_in_graph(forward_batch)
@@ -14,9 +14,12 @@ from sglang.srt.configs.model_config import (
get_minimax_sparse_score_type, get_minimax_sparse_score_type,
) )
from sglang.srt.environ import envs 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.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.server_args import m3_fp8_attn_gemm_enabled
from sglang.srt.utils import is_npu 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.sparse.init_forward_metadata_out_graph(forward_batch, in_capture)
self.dense.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): def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch):
self.sparse.init_forward_metadata_in_graph(forward_batch) self.sparse.init_forward_metadata_in_graph(forward_batch)
self.dense.init_forward_metadata_in_graph(forward_batch) self.dense.init_forward_metadata_in_graph(forward_batch)
@@ -4,11 +4,14 @@ from types import SimpleNamespace
from typing import TYPE_CHECKING, Callable, List, Optional from typing import TYPE_CHECKING, Callable, List, Optional
from sglang.srt.batch_overlap import two_batch_overlap 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: if TYPE_CHECKING:
from sglang.srt.layers.attention.verify_mask import VerifyMask 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): class TboAttnBackend(AttentionBackend):
@@ -117,6 +120,11 @@ class TboAttnBackend(AttentionBackend):
forward_batch=child_fb_view, in_capture=False 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): def init_forward_metadata_in_graph(self, forward_batch: ForwardBatch):
self.primary.init_forward_metadata_in_graph(forward_batch=forward_batch) self.primary.init_forward_metadata_in_graph(forward_batch=forward_batch)
if not self._children_use_cuda_graph(): if not self._children_use_cuda_graph():
@@ -21,7 +21,7 @@ from sglang.kernels.ops.kvcache.trtllm_mha_page_table import (
build_trtllm_mha_page_table, build_trtllm_mha_page_table,
) )
from sglang.srt.environ import envs 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 ( from sglang.srt.layers.attention.flashinfer_backend import (
FlashInferAttnBackend, FlashInferAttnBackend,
FlashInferMultiStepDraftBackend, FlashInferMultiStepDraftBackend,
@@ -103,11 +103,11 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend):
supports_ragged_verify_graph: bool = True 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. # Prefill metadata init snapshots all scheduler-shared inputs pre-replay.
if forward_mode == ForwardMode.EXTEND: if fm == ForwardMode.EXTEND:
return SharedReadBoundary.PRE_REPLAY return SharedReadEnds.PRE_REPLAY
return super().shared_read_boundary(forward_mode) return super().shared_read_ends(fm)
def __init__( def __init__(
self, self,
@@ -44,7 +44,10 @@ from sglang.srt.distributed.parallel_state import (
) )
from sglang.srt.dllm.config import DllmConfig from sglang.srt.dllm.config import DllmConfig
from sglang.srt.environ import envs 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.attention.dsa.utils import is_dsa_enable_prefill_cp
from sglang.srt.layers.dp_attention import ( from sglang.srt.layers.dp_attention import (
DpPaddingMode, DpPaddingMode,
@@ -439,33 +442,35 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
self.in_graph_metadata_prep_done = make_external_event(self.device_module) self.in_graph_metadata_prep_done = make_external_event(self.device_module)
event = self.in_graph_metadata_prep_done event = self.in_graph_metadata_prep_done
if event is not None: 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. # resolution below never hands out an unrecorded event.
event.record() event.record()
def _resolve_shared_read_boundary( def _replay_attn_backend(self) -> AttentionBackend:
self, attn_backend, forward_mode # Under pdmux each stream replays on its own group member.
) -> SharedReadBoundary: if self.enable_pdmux:
"""Where this replay records its shared-read-done event: the backend's return self.model_runner.decode_attn_backend_group[get_current_stream_idx()]
declaration, demoted when this runner cannot record at that point. return self.attn_backend
UNKNOWN records nothing (scheduler keeps the coarse fence)."""
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 forward_mode.is_target_verify():
if not self.model_runner.spec_algorithm.is_last_shared_read_phase( if not self.model_runner.spec_algorithm.is_last_shared_read_phase(
forward_mode forward_mode
): ):
return SharedReadBoundary.UNKNOWN return SharedReadEnds.UNKNOWN
elif not forward_mode.is_decode(): elif not forward_mode.is_decode():
return SharedReadBoundary.UNKNOWN return SharedReadEnds.UNKNOWN
boundary = attn_backend.shared_read_boundary(forward_mode) declared = attn_backend.shared_read_ends(forward_mode)
if ( if (
boundary is SharedReadBoundary.IN_REPLAY declared is SharedReadEnds.IN_REPLAY
and self.in_graph_metadata_prep_done is None and self.in_graph_metadata_prep_done is None
): ):
# TODO: PRE_REPLAY is EARLIER than the declared boundary; POST_REPLAY # TODO: this lands EARLIER than declared; POST_REPLAY is the sound one.
# is the sound demotion for a backend that really reads in-graph. return SharedReadEnds.PRE_REPLAY
return SharedReadBoundary.PRE_REPLAY return declared
return boundary
def _publish_read_done(self, in_graph: bool): def _publish_read_done(self, in_graph: bool):
"""Hand the scheduler's WAR barrier the event marking this phase's """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 and forward_batch.spec_info is not None
): ):
forward_batch.spec_info.custom_mask = buffers.custom_mask forward_batch.spec_info.custom_mask = buffers.custom_mask
if self.enable_pdmux:
stream_idx = get_current_stream_idx() attn_backend = self._replay_attn_backend()
attn_backend = self.model_runner.decode_attn_backend_group[stream_idx]
else:
attn_backend = self.attn_backend
fb_view = build_replay_fb_view( fb_view = build_replay_fb_view(
forward_batch=forward_batch, forward_batch=forward_batch,
buffers=buffers, buffers=buffers,
@@ -1316,8 +1318,8 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
timer_ctx = device_timer_ctx( timer_ctx = device_timer_ctx(
self.model_runner.device_timer, forward_batch.forward_mode.name.lower() self.model_runner.device_timer, forward_batch.forward_mode.name.lower()
) )
shared_read_boundary = self._resolve_shared_read_boundary( shared_read_ends = self._resolve_shared_read_ends(
self.attn_backend, forward_batch.forward_mode self._replay_attn_backend(), forward_batch.forward_mode
) )
with timer_ctx, self.backend.replay_session(): with timer_ctx, self.backend.replay_session():
self.load_batch(forward_batch, pp_proxy_tensors) self.load_batch(forward_batch, pp_proxy_tensors)
@@ -1335,15 +1337,15 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
else "" else ""
), ),
) )
if shared_read_boundary is SharedReadBoundary.PRE_REPLAY: if shared_read_ends is SharedReadEnds.PRE_REPLAY:
self._publish_read_done(in_graph=False) self._publish_read_done(in_graph=False)
output = self.backend.replay(self._replay_graph_key, forward_batch) 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) 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) self._publish_read_done(in_graph=False)
if isinstance(output, LogitsProcessorOutput): if isinstance(output, LogitsProcessorOutput):
@@ -6,7 +6,7 @@ from typing import Optional
import torch import torch
from sglang.srt.environ import envs 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.model_executor.forward_batch_info import ForwardMode
from sglang.srt.utils import is_cuda 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(): if not model_runner.spec_algorithm.is_none():
return return
# The record lands right after replay prep, so PRE_REPLAY only. # The record lands right after replay prep, so PRE_REPLAY only.
boundary = model_runner.attn_backend.shared_read_boundary( declared = model_runner.attn_backend.shared_read_ends(forward_batch.forward_mode)
forward_batch.forward_mode if declared is not SharedReadEnds.PRE_REPLAY:
)
if boundary is not SharedReadBoundary.PRE_REPLAY:
return return
logger.info_once( logger.info_once(
"Prefill shared-read-done fastpath active (%s)", "Prefill shared-read-done fastpath active (%s)",
@@ -1,135 +1,81 @@
import contextlib
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import create_autospec
import pytest import pytest
import torch
from sglang.srt.layers.attention.base_attn_backend import SharedReadBoundary from sglang.srt.layers.attention.base_attn_backend import (
from sglang.srt.model_executor.forward_batch_info import ForwardMode, PPProxyTensors AttentionBackend,
SharedReadEnds,
)
from sglang.srt.model_executor.forward_batch_info import ForwardMode
from sglang.srt.model_executor.runner.decode_cuda_graph_runner import ( from sglang.srt.model_executor.runner.decode_cuda_graph_runner import (
DecodeCudaGraphRunner, DecodeCudaGraphRunner,
) )
from sglang.srt.model_executor.runner.shape_key import ShapeKey
from sglang.test.ci.ci_register import register_cpu_ci from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=1, suite="base-a-test-cpu") register_cpu_ci(est_time=1, suite="base-a-test-cpu")
DECODE = ForwardMode.DECODE
class _SpecAlgorithm: VERIFY = ForwardMode.TARGET_VERIFY
def __init__(self, target_verify_war: bool = False): EXTEND = ForwardMode.EXTEND
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): def _runner(*, owns_verify: bool = False, has_marker: bool = False):
"""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 = DecodeCudaGraphRunner.__new__(DecodeCudaGraphRunner)
runner.model_runner = SimpleNamespace( runner.model_runner = SimpleNamespace(
spec_algorithm=_SpecAlgorithm(target_verify_war), spec_algorithm=SimpleNamespace(
device_timer=None, is_last_shared_read_phase=lambda fm: owns_verify and fm.is_target_verify()
is_draft_worker=False, ),
shared_read_done_event=None, shared_read_done_event=None,
) )
runner.in_graph_metadata_prep_done = object() if has_marker else None runner.in_graph_metadata_prep_done = object() if has_marker else None
return runner return runner
def test_unrelated_modes_never_publish(): def _backend(declared: SharedReadEnds):
# This runner owns the fence for decode / target verify only; every other # Spec'd against the real ABC so a rename fails here, not at runtime.
# mode stays on the coarse wait, even with a marker available. backend = create_autospec(AttentionBackend, instance=True)
assert ( backend.shared_read_ends.return_value = declared
_runner(has_marker=True)._resolve_shared_read_boundary( return backend
_attn_backend(), ForwardMode.EXTEND
)
is SharedReadBoundary.UNKNOWN
)
def test_post_replay_declaration_is_not_advanced(): @pytest.mark.parametrize(
# A backend that keeps reading shared state across the whole graph declares "mode, owns_verify, declared, has_marker, expected",
# POST_REPLAY. Having an in-graph marker must not pull the fence earlier. [
assert ( # Only decode / target verify publish; anything else keeps the coarse fence.
_runner(target_verify_war=True, has_marker=True)._resolve_shared_read_boundary( (EXTEND, False, SharedReadEnds.IN_REPLAY, True, SharedReadEnds.UNKNOWN),
_attn_backend(SharedReadBoundary.POST_REPLAY), ForwardMode.TARGET_VERIFY # Target verify publishes only when it is the step's last reading phase.
) (VERIFY, False, SharedReadEnds.IN_REPLAY, True, SharedReadEnds.UNKNOWN),
is SharedReadBoundary.POST_REPLAY (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): def test_publish_read_done():
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) runner = _runner(has_marker=True)
marker = runner.in_graph_metadata_prep_done recorded = []
runner.attn_backend = _attn_backend()
runner.device_module = SimpleNamespace( runner.device_module = SimpleNamespace(
Event=lambda: (_ for _ in ()).throw( Event=lambda: SimpleNamespace(record=lambda: recorded.append("record"))
AssertionError("execute must reuse the graph-recorded event")
)
) )
calls = []
forward_batch = _execute_harness(runner, calls)
result = runner.execute(forward_batch) runner._publish_read_done(in_graph=True)
# In-graph: hand over the graph-recorded marker, do not record a new event.
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 marker = runner.in_graph_metadata_prep_done
runner.attn_backend = _attn_backend() assert runner.model_runner.shared_read_done_event is marker
runner.device_module = SimpleNamespace(Event=lambda: None) assert recorded == []
runner.execute(_execute_harness(runner, [], ForwardMode.TARGET_VERIFY)) runner._publish_read_done(in_graph=False)
assert recorded == ["record"]
expected = marker if supported else None assert runner.model_runner.shared_read_done_event is not marker
assert runner.model_runner.shared_read_done_event is expected
if __name__ == "__main__": if __name__ == "__main__":
@@ -1,9 +1,13 @@
from types import SimpleNamespace from types import SimpleNamespace
from unittest.mock import create_autospec
import pytest import pytest
from sglang.srt.environ import envs 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.forward_batch_info import ForwardMode
from sglang.srt.model_executor.runner_utils import ( from sglang.srt.model_executor.runner_utils import (
maybe_publish_prefill_shared_read_done, maybe_publish_prefill_shared_read_done,
@@ -23,12 +27,13 @@ class _Event:
def _model_runner(*, spec_algorithm=SpeculativeAlgorithm.NONE, compliant=True): def _model_runner(*, spec_algorithm=SpeculativeAlgorithm.NONE, compliant=True):
boundary = ( declared = SharedReadEnds.PRE_REPLAY if compliant else SharedReadEnds.UNKNOWN
SharedReadBoundary.PRE_REPLAY if compliant else SharedReadBoundary.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( return SimpleNamespace(
spec_algorithm=spec_algorithm, spec_algorithm=spec_algorithm,
attn_backend=SimpleNamespace(shared_read_boundary=lambda mode: boundary), attn_backend=attn_backend,
shared_read_done_event=None, 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)), (_model_runner(), _batch(ForwardMode.DECODE)),
# The algorithm has a later prefill reader or unverified ownership. # The algorithm has a later prefill reader or unverified ownership.
(_model_runner(spec_algorithm=SpeculativeAlgorithm.EAGLE), _batch()), (_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()), (_model_runner(compliant=False), _batch()),
): ):
maybe_publish_prefill_shared_read_done(runner, batch, _DEVICE_MODULE) maybe_publish_prefill_shared_read_done(runner, batch, _DEVICE_MODULE)