[Refactor][TCPCG] Standardize diffusion eager graph wrappers (7/9)

python/sglang/multimodal_gen/runtime/layers/attention/layer.py

@@ -2101,16 +2101,15 @@ def _make_breakable_attention_forward(forward_method):
     disabled this is a transparent pass-through to the original method.
     """
 
-    def _forward_boxing_tuples(*args, **kwargs):
+    @eager_on_graph
+    def _eager_attention(*args, **kwargs):
         out = forward_method(*args, **kwargs)
         return _BCGBoxedTupleOutput(out) if isinstance(out, tuple) else out
 
-    bcg_forward = eager_on_graph(True)(_forward_boxing_tuples)
-
     @functools.wraps(forward_method)
     def forward(self, *args, **kwargs):
         if is_in_breakable_cuda_graph():
-            out = bcg_forward(self, *args, **kwargs)
+            out = _eager_attention(self, *args, **kwargs)
             return out.astuple() if isinstance(out, _BCGBoxedTupleOutput) else out
         return forward_method(self, *args, **kwargs)
 

python/sglang/multimodal_gen/runtime/models/dits/minimax_h3.py

@@ -775,7 +775,9 @@ def _minimax_h3_attention_core_impl(
     return out
 
 
-_minimax_h3_attention_core_bcg = eager_on_graph(True)(_minimax_h3_attention_core_impl)
+@eager_on_graph
+def _eager_attention_core(*args, **kwargs):
+    return _minimax_h3_attention_core_impl(*args, **kwargs)
 
 
 class MiniMaxH3Attention(nn.Module):
@@ -1173,7 +1175,7 @@ def forward(
             gate_compress = gate_flat.view(total, self.num_heads, self.head_dim)
 
         attention_core = (
-            _minimax_h3_attention_core_bcg
+            _eager_attention_core
             if self.bcg_breakpoint
             else _minimax_h3_attention_core_impl
         )
@@ -2395,8 +2397,8 @@ def build_rope_cache(
         self.release_mps_non_layer_weights("rope")
         return result
 
-    @eager_on_graph(True)
-    def _embed(
+    @eager_on_graph
+    def _eager_embed(
         self,
         *,
         x: torch.Tensor,
@@ -2428,7 +2430,7 @@ def _embed(
         elif torch.is_tensor(refined_prompt_embeds_length):
             # BCG turns this request-varying host constant into a scalar input
             # so different live lengths can replay one padded-text signature.
-            # _embed is an eager graph break, so this value is read outside
+            # _eager_embed is an eager graph break, so this value is read outside
             # captured CUDA graphs.
             text_len = int(refined_prompt_embeds_length.item())
         else:
@@ -2692,7 +2694,7 @@ def forward(self, **kwargs: Any) -> tuple[torch.Tensor, torch.Tensor]:
         audio_pos = audio_pos.to(device)
         text_pos = text_pos.to(device)
 
-        decoder_input, t_emb = self._embed(
+        decoder_input, t_emb = self._eager_embed(
             x=x,
             audio_x=audio_x,
             text_embeddings_selected=text_selected,

python/sglang/multimodal_gen/runtime/models/dits/minimax_h3_vdn_attention.py

@@ -140,7 +140,7 @@ def forward(
         beta = self.linear_attention.beta(x)
         gate_hidden, _ = self.linear_attention.output_gate.down(x)
         attention_core = (
-            _hybrid_attention_core_bcg
+            _eager_hybrid_attention_core
             if attention.bcg_breakpoint
             else _minimax_h3_hybrid_attention_core_impl
         )
@@ -563,9 +563,9 @@ def _vdn_return_to_rows(
     return merged[0], linear_rows
 
 
-_hybrid_attention_core_bcg = eager_on_graph(True)(
-    _minimax_h3_hybrid_attention_core_impl
-)
+@eager_on_graph
+def _eager_hybrid_attention_core(*args, **kwargs):
+    return _minimax_h3_hybrid_attention_core_impl(*args, **kwargs)
 
 
 def prepare_hybrid_attention_metadata(

python/sglang/srt/compilation/cuda_piecewise_backend.py

@@ -17,7 +17,9 @@
     get_pcg_capture_stream,
     is_in_torch_compile_warmup,
 )
-from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
+from sglang.srt.model_executor.runner_backend_utils.weak_ref_tensor import (
+    weak_ref_tensors,
+)
 from sglang.srt.model_executor.runner_utils.pool import (
     graph_pool_capture_scope,
     graph_pool_replay_scope,

python/sglang/srt/layers/attention/dsa/dsa_indexer.py

@@ -149,7 +149,7 @@
 
 if _is_cuda or _is_hip:
     # Plain-torch graph helpers: usable wherever the split-op surface is.
-    from sglang.srt.layers.attention.dsa.dsa_prefill_cuda_graph import (
+    from sglang.srt.layers.attention.dsa.head_gate import (
         logits_head_gate_graph,
         scale_head_gate_graph,
     )

python/sglang/srt/layers/attention/dsa/dsa_prefill_cuda_graph.py

@@ -14,11 +14,10 @@
     get_tc_piecewise_forward_context,
     is_in_tc_piecewise_cuda_graph,
 )
-from sglang.srt.utils import is_cuda, is_hip
+from sglang.srt.utils import is_cuda
 from sglang.srt.utils.custom_op import register_custom_op
 
 _is_cuda = is_cuda()
-_is_hip = is_hip()
 
 GRAPH_WEIGHTS_PROJ_LORA_ERROR = (
     "DSA indexer weights_proj LoRA is incompatible with "
@@ -32,60 +31,6 @@ def _is_in_piecewise_or_breakable_cuda_graph() -> bool:
     return is_in_tc_piecewise_cuda_graph() or is_in_breakable_cuda_graph()
 
 
-if _is_cuda or _is_hip:
-
-    def _scale_head_gate_graph_fake_impl(
-        weights_raw: torch.Tensor,
-        n_heads_inv_sqrt: float,
-        softmax_scale: float,
-        q_scale: torch.Tensor,
-    ) -> torch.Tensor:
-        return torch.empty(
-            (weights_raw.shape[0], weights_raw.shape[1], q_scale.shape[-1]),
-            dtype=torch.float32,
-            device=weights_raw.device,
-        )
-
-    # In-graph (PCG/BCG) head gate for the fused path: weights_proj is folded
-    # into wk_weights_proj, so weights_raw is precomputed and there is no GEMM.
-    @register_custom_op(fake_impl=_scale_head_gate_graph_fake_impl)
-    def scale_head_gate_graph(
-        weights_raw: torch.Tensor,
-        n_heads_inv_sqrt: float,
-        softmax_scale: float,
-        q_scale: torch.Tensor,
-    ) -> torch.Tensor:
-        weights = weights_raw * n_heads_inv_sqrt
-        return weights.unsqueeze(-1) * q_scale * softmax_scale
-
-    def _logits_head_gate_graph_fake_impl(
-        x: torch.Tensor,
-        weight: torch.Tensor,
-        n_heads_inv_sqrt: float,
-        softmax_scale: float,
-        q_scale: torch.Tensor,
-    ) -> torch.Tensor:
-        return torch.empty(
-            (x.shape[0], weight.shape[0], q_scale.shape[-1]),
-            dtype=torch.float32,
-            device=x.device,
-        )
-
-    # In-graph (PCG/BCG) head gate for the NON-prefill path
-    @register_custom_op(fake_impl=_logits_head_gate_graph_fake_impl)
-    def logits_head_gate_graph(
-        x: torch.Tensor,
-        weight: torch.Tensor,
-        n_heads_inv_sqrt: float,
-        softmax_scale: float,
-        q_scale: torch.Tensor,
-    ) -> torch.Tensor:
-        out = torch.mm(x, weight.t(), out_dtype=torch.float32)
-        weights = out * n_heads_inv_sqrt
-        weights = weights.unsqueeze(-1) * q_scale * softmax_scale
-        return weights
-
-
 @register_custom_op(mutates_args=["topk_result"])
 @register_split_op()
 def pcg_dsa_indexer_prefill_split(

python/sglang/srt/layers/attention/dsa/head_gate.py

@@ -0,0 +1,63 @@
+from __future__ import annotations
+
+import torch
+
+from sglang.srt.utils import is_cuda, is_hip
+from sglang.srt.utils.custom_op import register_custom_op
+
+_is_cuda = is_cuda()
+_is_hip = is_hip()
+
+
+if _is_cuda or _is_hip:
+
+    def _scale_head_gate_graph_fake_impl(
+        weights_raw: torch.Tensor,
+        n_heads_inv_sqrt: float,
+        softmax_scale: float,
+        q_scale: torch.Tensor,
+    ) -> torch.Tensor:
+        return torch.empty(
+            (weights_raw.shape[0], weights_raw.shape[1], q_scale.shape[-1]),
+            dtype=torch.float32,
+            device=weights_raw.device,
+        )
+
+    # In-graph head gate for the fused path: weights_proj is folded
+    # into wk_weights_proj, so weights_raw is precomputed and there is no GEMM.
+    @register_custom_op(fake_impl=_scale_head_gate_graph_fake_impl)
+    def scale_head_gate_graph(
+        weights_raw: torch.Tensor,
+        n_heads_inv_sqrt: float,
+        softmax_scale: float,
+        q_scale: torch.Tensor,
+    ) -> torch.Tensor:
+        weights = weights_raw * n_heads_inv_sqrt
+        return weights.unsqueeze(-1) * q_scale * softmax_scale
+
+    def _logits_head_gate_graph_fake_impl(
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        n_heads_inv_sqrt: float,
+        softmax_scale: float,
+        q_scale: torch.Tensor,
+    ) -> torch.Tensor:
+        return torch.empty(
+            (x.shape[0], weight.shape[0], q_scale.shape[-1]),
+            dtype=torch.float32,
+            device=x.device,
+        )
+
+    # In-graph head gate for the NON-prefill path
+    @register_custom_op(fake_impl=_logits_head_gate_graph_fake_impl)
+    def logits_head_gate_graph(
+        x: torch.Tensor,
+        weight: torch.Tensor,
+        n_heads_inv_sqrt: float,
+        softmax_scale: float,
+        q_scale: torch.Tensor,
+    ) -> torch.Tensor:
+        out = torch.mm(x, weight.t(), out_dtype=torch.float32)
+        weights = out * n_heads_inv_sqrt
+        weights = weights.unsqueeze(-1) * q_scale * softmax_scale
+        return weights

python/sglang/srt/model_executor/runner_backend/breakable_cuda_graph_backend.py

@@ -126,9 +126,7 @@ def capture_one(
                 post_warmup_hook()
 
         graph = BreakableCUDAGraph(self.deduped_cuda_graph)
-        captured_fn = (
-            eager_on_graph(True)(forward_fn) if self._debug_eager else forward_fn
-        )
+        captured_fn = eager_on_graph(forward_fn) if self._debug_eager else forward_fn
         size = shape_key.size
         if self._shared_output_buffer is None:
             capacity_rows = self._cuda_graph_runner.cuda_graph_output_capacity_rows(
@@ -266,7 +264,7 @@ def replay(
         **kwargs,
     ) -> Any:
         with graph_pool_replay_scope():
-            self._graphs[shape_key].replay()
+            self._graphs[shape_key].replay(static_forward_batch)
         return self._outputs[shape_key]
 
     def cleanup(self) -> None:

python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py

@@ -22,6 +22,8 @@
 buffers to keep break-point tensors at stable addresses.
 """
 
+import functools
+import inspect
 import threading
 import warnings
 from contextvars import ContextVar
@@ -166,7 +168,9 @@ def _weak_ref_if_tensor(x):
     if torch.is_tensor(x):
         if x.numel() == 0 or x.device.type == "cpu":
             return x
-        from sglang.srt.compilation.weak_ref_tensor import weak_ref_tensors
+        from sglang.srt.model_executor.runner_backend_utils.weak_ref_tensor import (
+            weak_ref_tensors,
+        )
 
         return weak_ref_tensors(x)
     if isinstance(x, tuple):
@@ -216,11 +220,30 @@ def _copy_output(dst: Any, src: Any) -> Any:
     return src
 
 
-def eager_on_graph(enable: bool, capture_stub: Optional[Callable] = None):
+def eager_on_graph(
+    fn: Optional[Callable] = None,
+    capture_stub: Optional[Callable] = None,
+    *,
+    enable: Optional[bool] = None,
+):
+    """Record an eager call between captured segments.
+
+    A named ``forward_batch`` argument is rebound to the prepared serving batch
+    on replay. All other arguments retain their capture-time identity; tensors
+    must therefore use the static buffers owned by the graph runner.
+    """
+
+    # Transitional support while model callers migrate to the bare decorator.
+    if isinstance(fn, bool):
+        enable, fn = fn, None
+    if enable is False:
+        return lambda inner: inner
+
     def decorator(inner: Callable):
-        if not enable:
-            return inner
+        signature = inspect.signature(inner)
+        has_batch = "forward_batch" in signature.parameters
 
+        @functools.wraps(inner)
         def wrapper(*args, **kwargs):
             capture = _current_capture_var.get()
             if capture is None:
@@ -256,9 +279,31 @@ def wrapper(*args, **kwargs):
             # captured segment. Keep a strong reference so replay can safely
             # copy fresh eager output into that bridge buffer.
             captured_output = output
+            replay_args = (
+                signature.bind(*captured_args, **captured_kwargs) if has_batch else None
+            )
 
-            def replay_fn():
-                new_out = captured_inner(*captured_args, **captured_kwargs)
+            if replay_args is not None:
+                # Batches own request metadata and tensors. Keep only static tensor
+                # operands between calls, never a captured or previous serving batch.
+                replay_args.arguments["forward_batch"] = None
+                captured_args, captured_kwargs = (), {}
+
+            def replay_fn(forward_batch):
+                if has_batch:
+                    if forward_batch is None:
+                        raise ValueError(
+                            "This eager region requires a replay ForwardBatch"
+                        )
+                    replay_args.arguments["forward_batch"] = forward_batch
+                    try:
+                        new_out = captured_inner(
+                            *replay_args.args, **replay_args.kwargs
+                        )
+                    finally:
+                        replay_args.arguments["forward_batch"] = None
+                else:
+                    new_out = captured_inner(*captured_args, **captured_kwargs)
                 return _copy_output(captured_output, new_out)
 
             capture.cuda_graph._break_fns.append(replay_fn)
@@ -269,7 +314,7 @@ def replay_fn():
 
         return wrapper
 
-    return decorator
+    return decorator(fn) if fn is not None else decorator
 
 
 class BreakableCUDAGraph:
@@ -278,17 +323,17 @@ class BreakableCUDAGraph:
 
     def __init__(self, deduped_cuda_graph=None) -> None:
         self._segments: list[Any] = []
-        self._break_fns: list[Callable[[], Any]] = []
+        self._break_fns: list[Callable[[Any], Any]] = []
         self._deduped_cuda_graph = deduped_cuda_graph
 
-    def replay(self) -> None:
+    def replay(self, forward_batch=None) -> None:
         stream = get_device_module().current_stream()
         token = _current_stream_var.set(stream)
         try:
             for i, seg in enumerate(self._segments):
                 seg.replay()
                 if i < len(self._break_fns):
-                    self._break_fns[i]()
+                    self._break_fns[i](forward_batch)
         finally:
             _current_stream_var.reset(token)
 
@@ -408,8 +453,7 @@ def _end_current_segment(self) -> None:
         self._current_graph_needs_instantiate = False
 
 
-@eager_on_graph(True)
+@eager_on_graph
 def break_graph() -> None:
     """Insert a graph break. The @eager_on_graph decorator does the actual
     segment split; this function body intentionally does nothing."""
-    pass