[Refactor][TCPCG] Bind live forward batches in BCG eager calls (2/9)

python/sglang/srt/model_executor/runner/prefill_cuda_graph_runner.py

@@ -124,6 +124,7 @@
     PrefillInputBuffers,
 )
 from sglang.srt.model_executor.runner_utils.capture_mode import model_capture_mode
+from sglang.srt.model_executor.runner_utils.forward_batch import set_forward_batch
 from sglang.srt.model_executor.runner_utils.pool import (
     get_or_create_global_graph_capture_stream,
 )
@@ -747,6 +748,7 @@ def _prefill_forward_context(
         raw_num_tokens: Optional[int] = None,
     ):
         with (
+            set_forward_batch(forward_batch),
             forward_context(
                 ForwardContext(attn_backend=self.model_runner.attn_backend)
             ),
@@ -879,6 +881,7 @@ def run_dummy_multimodal_deepstack_forward(
         set_is_extend_in_batch(False)
 
         with (
+            set_forward_batch(fb),
             forward_context(
                 ForwardContext(attn_backend=self.model_runner.attn_backend)
             ),

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

@@ -127,9 +127,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(

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

@@ -22,6 +22,7 @@
 buffers to keep break-point tensors at stable addresses.
 """
 
+import functools
 import threading
 import warnings
 from contextvars import ContextVar
@@ -218,11 +219,27 @@ def _copy_output(dst: Any, src: Any) -> Any:
     return src
 
 
-def eager_on_graph(enable: bool, capture_stub: Optional[Callable] = None):
-    def decorator(inner: Callable):
-        if not enable:
-            return inner
+def eager_on_graph(
+    fn: Optional[Callable] = None,
+    capture_stub: Optional[Callable] = None,
+    *,
+    enable: Optional[bool] = None,
+):
+    """Record an eager call between captured segments.
+
+    Arguments retain their capture-time identity. Tensors must use the static
+    buffers owned by the graph runner; request-specific state must be read
+    inside the eager function at execution time.
+    """
 
+    # 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):
+        @functools.wraps(inner)
         def wrapper(*args, **kwargs):
             capture = _current_capture_var.get()
             if capture is None:
@@ -271,7 +288,7 @@ def replay_fn():
 
         return wrapper
 
-    return decorator
+    return decorator(fn) if fn is not None else decorator
 
 
 class BreakableCUDAGraph:
@@ -410,8 +427,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

python/sglang/srt/model_executor/runner_utils/forward_batch.py

@@ -0,0 +1,47 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Current execution batch for eager regions in prefill graphs.
+
+The prefill runner owns this scope. Eager callbacks resolve the batch when
+executed instead of retaining capture-time request metadata. This state is
+independent of ForwardContext, which selects the attention backend.
+"""
+
+from contextlib import contextmanager
+from contextvars import ContextVar
+from typing import TYPE_CHECKING, Iterator
+
+if TYPE_CHECKING:
+    from sglang.srt.model_executor.forward_batch_info import ForwardBatch
+
+_current_batch: ContextVar["ForwardBatch | None"] = ContextVar(
+    "forward_batch", default=None
+)
+
+
+def get_forward_batch() -> "ForwardBatch":
+    batch = _current_batch.get()
+    if batch is None:
+        raise RuntimeError("No forward batch is set for this execution")
+    return batch
+
+
+@contextmanager
+def set_forward_batch(forward_batch: "ForwardBatch") -> Iterator[None]:
+    """Publish the current batch and restore the previous one on scope exit."""
+    reset_token = _current_batch.set(forward_batch)
+    try:
+        yield
+    finally:
+        _current_batch.reset(reset_token)

test/registered/unit/model_executor/test_eager_on_graph.py

@@ -0,0 +1,174 @@
+"""CPU tests of scoped eager replay without requiring CUDA graph capture."""
+
+import sys
+from contextlib import contextmanager
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import (
+    breakable_cuda_graph as bcg,
+)
+from sglang.srt.model_executor.runner_utils.forward_batch import (
+    get_forward_batch,
+    set_forward_batch,
+)
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=3, suite="base-a-test-cpu")
+
+
+@contextmanager
+def recording_capture():
+    events = []
+    graph = bcg.BreakableCUDAGraph()
+    capture = SimpleNamespace(
+        cuda_graph=graph,
+        _barrier_fn=None,
+        _end_current_segment=lambda: events.append("end"),
+        _begin_new_segment=lambda: events.append("begin"),
+    )
+    token = bcg._current_capture_var.set(capture)
+    try:
+        yield graph, events
+    finally:
+        bcg._current_capture_var.reset(token)
+
+
+def test_method_replay_reads_current_batch_and_preserves_output_address(monkeypatch):
+    class Layer:
+        @bcg.eager_on_graph
+        def _eager_scale(self, x):
+            return x * get_forward_batch().scale
+
+    layer = Layer()
+    x = torch.tensor([2.0])
+    with set_forward_batch(SimpleNamespace(scale=3)):
+        with recording_capture() as (graph, events):
+            output = layer._eager_scale(x)
+    assert events == ["end", "begin"]
+    pointer = output.data_ptr()
+    # Exercise the public, argument-free replay API between two segments.
+    graph._segments = [
+        SimpleNamespace(replay=lambda: events.append("first")),
+        SimpleNamespace(replay=lambda: events.append("last")),
+    ]
+    monkeypatch.setattr(
+        bcg, "get_device_module", lambda: SimpleNamespace(current_stream=lambda: None)
+    )
+    x.fill_(4)
+    with set_forward_batch(SimpleNamespace(scale=5)):
+        graph.replay()
+    assert events == ["end", "begin", "first", "last"]
+    assert output.item() == 20 and output.data_ptr() == pointer
+    with pytest.raises(RuntimeError, match="No forward batch"):
+        graph.replay()
+
+
+def test_method_capture_stub_and_real_replay():
+    calls = []
+
+    class Layer:
+        def _capture_stub(self, x):
+            calls.append("stub")
+            return torch.zeros_like(x)
+
+        @bcg.eager_on_graph(capture_stub=_capture_stub)
+        def _eager_scale(self, x):
+            calls.append("real")
+            return x * get_forward_batch().scale
+
+    layer = Layer()
+    with recording_capture() as (graph, _):
+        output = layer._eager_scale(torch.ones(2))
+    assert calls == ["stub"]
+    with set_forward_batch(SimpleNamespace(scale=7)):
+        graph._break_fns[0]()
+    assert calls == ["stub", "real"]
+    torch.testing.assert_close(output, torch.full((2,), 7.0))
+
+
+def test_plain_callable_and_eager_passthrough():
+    @bcg.eager_on_graph
+    def increment(x):
+        return x + 1
+
+    assert increment(torch.tensor(2)).item() == 3
+    x = torch.tensor(3)
+    with recording_capture() as (graph, _):
+        output = increment(x)
+    x.fill_(9)
+    graph._break_fns[0]()
+    assert output.item() == 10
+
+
+def test_replay_does_not_retain_capture_or_serving_batches():
+    import gc
+    import weakref
+
+    class Batch:
+        scale = 3
+
+    @bcg.eager_on_graph
+    def scale(x):
+        return x * get_forward_batch().scale
+
+    batch = Batch()
+    capture_ref = weakref.ref(batch)
+    with set_forward_batch(batch):
+        with recording_capture() as (graph, _):
+            scale(torch.ones(2))
+    del batch
+    gc.collect()
+    assert capture_ref() is None
+
+    batch = Batch()
+    live_ref = weakref.ref(batch)
+    with set_forward_batch(batch):
+        graph._break_fns[0]()
+    del batch
+    gc.collect()
+    assert live_ref() is None
+
+
+@pytest.mark.parametrize(
+    "decorator", [bcg.eager_on_graph(True), bcg.eager_on_graph(enable=True)]
+)
+def test_legacy_decorator_replays(decorator):
+    increment = decorator(lambda x: x + 1)
+    with recording_capture() as (graph, _):
+        output = increment(torch.tensor(2))
+    graph._break_fns[0]()
+    assert output.item() == 3
+
+
+def test_legacy_disabled_decorator_is_identity():
+    fn = lambda x: x + 1
+    assert bcg.eager_on_graph(False)(fn) is fn
+    assert bcg.eager_on_graph(enable=False)(fn) is fn
+
+
+def test_batch_scope_restores_outer_batch_after_exception():
+    outer = SimpleNamespace(scale=2)
+    with set_forward_batch(outer):
+        with pytest.raises(ValueError, match="eager failure"):
+            with set_forward_batch(SimpleNamespace(scale=9)):
+                assert get_forward_batch().scale == 9
+                raise ValueError("eager failure")
+        assert get_forward_batch() is outer
+    with pytest.raises(RuntimeError, match="No forward batch"):
+        get_forward_batch()
+
+
+def test_batch_scope_is_isolated_between_execution_contexts():
+    from contextvars import Context
+
+    with set_forward_batch(SimpleNamespace(scale=2)):
+        with pytest.raises(RuntimeError, match="No forward batch"):
+            Context().run(get_forward_batch)
+        assert get_forward_batch().scale == 2
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__, "-v"]))