[misc] Rename shared-read boundary to shared-read ends and fix wrapper delegation (#34982)
This commit is contained in:
Liangsheng Yin
+47
-101
@@ -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__":
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user