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

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

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,

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

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)

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)

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():

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,

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):

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)",