[misc] Rename the WAR read-done fastpath to shared-read-done (#34916)
This commit is contained in:
Liangsheng Yin
+136
@@ -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"]))
|
||||
@@ -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"]))
|
||||
+10
-8
@@ -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__":
|
||||
Reference in New Issue
Block a user