[misc] Rename the WAR read-done fastpath to shared-read-done (#34916)

This commit is contained in:
Liangsheng Yin
2026-08-15 15:02:02 -07:00
committed by GitHub
parent 4d0c5a89af
commit 0f7aaceda5
16 changed files with 224 additions and 249 deletions
@@ -21,7 +21,7 @@ if TYPE_CHECKING:
class SharedReadBoundary(Enum):
"""Where a backend's scheduler-shared reads end, relative to the replay;
the WAR read-done record must land at or after this point. IN_REPLAY
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()
+3 -3
View File
@@ -1709,9 +1709,9 @@ class Scheduler(
# (forceable via SGLANG_FORCE_COARSE_WAR_BARRIER).
if not self._war_barrier_enabled:
return
runner = self.model_worker.war_fastpath_runner
ev = runner.war_fastpath_read_done_event
runner.war_fastpath_read_done_event = None
runner = self.model_worker.last_shared_read_runner
ev = runner.shared_read_done_event
runner.shared_read_done_event = None
if ev is not None and not envs.SGLANG_FORCE_COARSE_WAR_BARRIER.get():
self.schedule_stream.wait_event(ev)
else:
+1 -1
View File
@@ -82,7 +82,7 @@ class BaseTpWorker(ABC):
pass
@property
def war_fastpath_runner(self):
def last_shared_read_runner(self):
# The runner that runs the step's LAST shared-buffer-reading phase --
# it owns the read-done event the scheduler's WAR barrier waits on.
# For a plain worker that's its own runner.
@@ -166,7 +166,6 @@ 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,
@@ -404,14 +403,10 @@ class ModelRunner:
# Init forward stream for overlap schedule
self.forward_stream = torch.get_device_module(self.device).Stream()
# WAR fast-path: a decode-graph forward publishes a fresh event here after
# 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)
)
# Published by the step's last shared-buffer-reading phase (decode graph,
# eagle draft extend, or prefill); the scheduler's WAR barrier waits on it
# then clears it. None -> coarse whole-forward wait_stream.
self.shared_read_done_event: Optional[torch.cuda.Event] = None
# CPU offload
set_offloader(
@@ -93,6 +93,7 @@ from sglang.srt.model_executor.runner_utils.capture_mode import (
from sglang.srt.model_executor.runner_utils.deepep_adapter import (
DeepEPCudaGraphRunnerAdapter,
)
from sglang.srt.model_executor.runner_utils.shared_read_event import make_external_event
from sglang.srt.multiplex.pdmux_context import get_current_stream_idx, get_stream_groups
from sglang.srt.runtime_context import get_flags, get_parallel, get_spec
from sglang.srt.speculative.ragged_verify import resolve_ragged_verify_layout
@@ -208,7 +209,10 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
speculative_num_draft_tokens: Optional[int] = None,
):
super().__init__(model_runner)
self._war_read_done_node_planted = False
# In-graph metadata prep: shared buffers -> in-graph private data
self.in_graph_metadata_prep_done: Optional[torch.cuda.Event] = None
# --- core state ------------------------------------------------
self.enable_torch_compile = get_flags().capture.enable_torch_compile
self.disable_padding = model_runner.server_args.disable_cuda_graph_padding
@@ -424,47 +428,55 @@ 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 _record_in_graph_metadata_prep_done(self):
# Purely a marker at this point in the graph; where the shared reads
# actually end is the attn backend's call.
if not torch.cuda.is_current_stream_capturing():
# Warmup shares this body. Breakable capture still plants: it opens
# segment 1 on context entry and every segment re-arms the node.
return
if self.in_graph_metadata_prep_done is None:
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
# resolution below never hands out an unrecorded event.
event.record()
def _war_read_done_record(self, attn_backend, forward_mode) -> SharedReadBoundary:
"""Where this replay records its WAR read-done event; UNKNOWN records
nothing."""
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)."""
if forward_mode.is_target_verify():
if not self.model_runner.spec_algorithm.is_war_publish_phase(forward_mode):
if not self.model_runner.spec_algorithm.is_last_shared_read_phase(
forward_mode
):
return SharedReadBoundary.UNKNOWN
elif not forward_mode.is_decode():
return SharedReadBoundary.UNKNOWN
boundary = attn_backend.shared_read_boundary(forward_mode)
if (
boundary is SharedReadBoundary.IN_REPLAY
and not self._war_read_done_node_planted
and self.in_graph_metadata_prep_done is None
):
# Non-capturing runs / no external-event support.
# 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
def _publish_war_read_done(self, in_graph: bool):
"""Publish the read-done event the scheduler's WAR barrier waits on."""
def _publish_read_done(self, in_graph: bool):
"""Hand the scheduler's WAR barrier the event marking this phase's
shared-buffer reads as done."""
if in_graph:
self.model_runner.war_fastpath_read_done_event = (
self.model_runner.war_read_done_event
)
# Reads end at the in-graph marker: wire it through, don't re-record.
self.model_runner.shared_read_done_event = self.in_graph_metadata_prep_done
else:
read_done = self.device_module.Event()
read_done.record()
self.model_runner.war_fastpath_read_done_event = read_done
self.model_runner.shared_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})
@@ -1066,7 +1078,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()
self._record_in_graph_metadata_prep_done()
# No invalidate_loc_cache() here: the unified pool translates its
# locs in `init_forward_metadata_out_graph`, so no cache to invalidate.
@@ -1304,7 +1316,7 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
timer_ctx = device_timer_ctx(
self.model_runner.device_timer, forward_batch.forward_mode.name.lower()
)
war_record = self._war_read_done_record(
shared_read_boundary = self._resolve_shared_read_boundary(
self.attn_backend, forward_batch.forward_mode
)
with timer_ctx, self.backend.replay_session():
@@ -1323,13 +1335,16 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
else ""
),
)
if war_record is SharedReadBoundary.PRE_REPLAY:
self._publish_war_read_done(in_graph=False)
if shared_read_boundary is SharedReadBoundary.PRE_REPLAY:
self._publish_read_done(in_graph=False)
output = self.backend.replay(self._replay_graph_key, forward_batch)
if war_record is SharedReadBoundary.POST_REPLAY:
self._publish_war_read_done(in_graph=False)
elif war_record is SharedReadBoundary.IN_REPLAY:
self._publish_war_read_done(in_graph=True)
if shared_read_boundary is SharedReadBoundary.IN_REPLAY:
self._publish_read_done(in_graph=True)
if shared_read_boundary is SharedReadBoundary.POST_REPLAY:
self._publish_read_done(in_graph=False)
if isinstance(output, LogitsProcessorOutput):
if self.is_dllm:
@@ -109,7 +109,9 @@ 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 import (
maybe_publish_prefill_shared_read_done,
)
from sglang.srt.model_executor.runner_utils.buffers import (
PrefillInputBuffers,
)
@@ -1758,7 +1760,7 @@ class PrefillCudaGraphRunner(BaseCudaGraphRunner):
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(
maybe_publish_prefill_shared_read_done(
self.model_runner, forward_batch, self.device_module
)
@@ -26,7 +26,6 @@ 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
make_war_read_done_event,
maybe_publish_prefill_war_read_done,
from sglang.srt.model_executor.runner_utils.shared_read_event import ( # noqa: F401
maybe_publish_prefill_shared_read_done,
)
@@ -1,4 +1,4 @@
"""WAR read-done event utilities for CUDA graph runners."""
"""Shared-read-done event utilities for CUDA graph runners."""
import logging
from typing import Optional
@@ -13,8 +13,8 @@ from sglang.srt.utils import is_cuda
logger = logging.getLogger(__name__)
def make_war_read_done_event(device_module) -> Optional[torch.cuda.Event]:
"""Create a persistent external event for CUDA graph capture."""
def make_external_event(device_module) -> Optional[torch.cuda.Event]:
"""Create a persistent external event, e.g., for CUDA graph capture."""
if not is_cuda():
return None
try:
@@ -23,7 +23,7 @@ def make_war_read_done_event(device_module) -> Optional[torch.cuda.Event]:
return None
def maybe_publish_prefill_war_read_done(
def maybe_publish_prefill_shared_read_done(
model_runner, forward_batch, device_module
) -> None:
"""Publish prefill read-done after compliant metadata initialization."""
@@ -42,9 +42,9 @@ def maybe_publish_prefill_war_read_done(
if boundary is not SharedReadBoundary.PRE_REPLAY:
return
logger.info_once(
"Prefill WAR read-done fastpath active (%s)",
"Prefill shared-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
model_runner.shared_read_done_event = read_done
@@ -212,7 +212,7 @@ class BaseSpecWorker(ABC):
return self.draft_worker.weight_load_time
@property
def war_fastpath_runner(self):
def last_shared_read_runner(self):
# The runner that runs the step's LAST shared-buffer-reading phase --
# it owns the read-done event the scheduler's WAR barrier waits on.
# Default is the target runner; override if the last phase runs
@@ -591,10 +591,10 @@ class EAGLEDraftExtendCudaGraphRunner(DecodeCudaGraphRunner):
# Snapshot built -- the forward is done reading the shared pool. Publish
# a read-done event the scheduler's WAR barrier waits on (draft extend
# is the EAGLE-family war-publish phase; last write wins the mailbox).
# is the EAGLE-family last shared-read phase; last write wins the mailbox).
read_done = self.device_module.Event()
read_done.record()
self.model_runner.war_fastpath_read_done_event = read_done
self.model_runner.shared_read_done_event = read_done
self.raw_bs = raw_bs
self.bs = bs
@@ -1055,7 +1055,7 @@ class EAGLEWorkerV2(BaseSpecWorker):
self.plan_stream, self.plan_stream_ctx = get_plan_stream(self.device)
@property
def war_fastpath_runner(self):
def last_shared_read_runner(self):
# Per the base contract: the step's last shared-buffer-reading phase is
# draft_extend, which runs on the draft runner.
return self._draft_worker.draft_runner
+2 -2
View File
@@ -127,8 +127,8 @@ class SpeculativeAlgorithm(Enum):
def supports_target_verify_for_draft(self) -> bool:
return self.is_dflash_family()
def is_war_publish_phase(self, forward_mode) -> bool:
# The step's last shared-buffer-reading phase owns the WAR read-done
def is_last_shared_read_phase(self, forward_mode) -> bool:
# The step's last shared-buffer-reading phase owns the shared-read-done
# publish. DSPARK has no draft_extend: its draft samples inside the
# verify graph, so verify is that last phase.
if self.is_dflash_family() or self.is_dspark():
@@ -92,8 +92,8 @@ class CustomSpecAlgo:
def supports_target_verify_for_draft(self) -> bool:
return False
def is_war_publish_phase(self, forward_mode) -> bool:
# The step's last shared-buffer-reading phase owns the WAR read-done publish.
def is_last_shared_read_phase(self, forward_mode) -> bool:
# The step's last shared-buffer-reading phase owns the shared-read-done publish.
return forward_mode.is_draft_extend_v2()
def supports_ragged_verify(self) -> bool: