[Refactor] bcg: Record graph breaks at capture without running them (#46719)

docs/docs/advanced_features/breakable_cuda_graph.mdx

@@ -30,15 +30,20 @@ This mode is intended for debugging only — it eliminates the performance benef
 
 ### Selective Graph Breaks in Model Code
 
-For production use, you can mark specific functions as "non-graphable" using the `@eager_on_graph` decorator. During CUDA graph capture, these functions run eagerly between captured graph segments. Outside of capture, they behave normally.
+For production use, you can mark specific functions as "non-graphable" using the `@eager_on_graph` decorator. During replay, these functions run eagerly between captured graph segments. Outside a breakable CUDA graph, they run as plain calls.
 
 ```python
 from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph import eager_on_graph
 
 @eager_on_graph
-def my_dynamic_op(x):
+def my_dynamic_op(x, out):
     # This op is incompatible with CUDA graph capture
-    return some_dynamic_operation(x)
+    out.copy_(some_dynamic_operation(x))
+
+def forward(self, x):
+    out = torch.empty_like(x)
+    my_dynamic_op(x, out)
+    return out
 ```
 
 You can also insert a bare graph break (no computation) using the `break_graph()` helper:
@@ -103,9 +108,9 @@ Begin capture (segment 1)
   ... graphable ops ...
   @eager_on_graph function encountered:
     1. End current capture segment
-    2. Run the function eagerly (allocates output tensors)
-    3. Record the function for later replay
-    4. Begin new capture segment
+    2. Record the function for later replay without running it; an optional
+       capture_stub runs host setup that later segments read
+    3. Begin new capture segment
   ... more graphable ops ...
 End capture (segment N)
 ```
@@ -123,15 +128,11 @@ For each segment i:
 Launch final CUDA graph segment
 ```
 
-The non-graph functions are re-invoked with the same tensor references as capture time. Since these references point to the CUDA graph's static input/output buffers, they see updated values on each replay.
+The non-graph functions are invoked with the same tensor references as capture time. Since these references point to the CUDA graph's static input/output buffers, they see updated values on each replay.
 
-### Output Writeback
+### Outputs
 
-When a non-graph function produces output during replay, the result must be written back into the same tensor buffers that downstream graph segments reference. The mechanism handles:
-
-- **Plain tensors**: In-place `copy_()` into the original buffer.
-- **Structured outputs** (dataclasses, objects with tensor attributes): Tensor fields are copied in-place; non-tensor fields are replaced.
-- **Dicts of tensors**: Tensor values are copied in-place; non-tensor values are replaced.
+A non-graph function writes its results into buffers its caller allocated before the break, so the downstream graph segments read the same addresses at every replay. Host-side results go into a holder object the caller owns. A function that returns a value fails with an assertion. Every captured shape is warmed up before capture, so lazy setup (JIT compilation, buffers sized by the first run) happens outside capture.
 
 ### Stream Fork/Join Tracking
 

python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py

@@ -556,9 +556,13 @@ class BaseBreakableCudaGraphRunner:
 
         # Warm up on the capture stream so cuBLAS/cuDNN/Triton workspaces and
         # any lazy JIT are materialized before capture (mirrors the LLM runner
-        # and torch.cuda.make_graphed_callables).
+        # and torch.cuda.make_graphed_callables). The BCG flag makes the warmup
+        # take the branches the capture records.
         self.device_module.synchronize()
-        with self.device_module.stream(self._capture_stream):
+        with (
+            self.device_module.stream(self._capture_stream),
+            enable_breakable_cuda_graph(),
+        ):
             for _ in range(2):
                 self.transformer(**static_kwargs)
         self._capture_stream.synchronize()

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

@@ -2441,7 +2441,7 @@ def _make_breakable_attention_forward(forward_method):
     """Wrap a DiT attention module's ``forward`` so it becomes a breakable
     CUDA graph (BCG) break point.
 
-    During BCG capture the whole attention forward runs eagerly between
+    At every BCG replay the whole attention forward runs eagerly between
     captured graph segments -- the sequence-parallel all-to-all collectives,
     varlen packing, and dynamic/sparse attention kernels that live here
     cannot (or should not) be captured into a static CUDA graph. The outputs

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

@@ -1239,7 +1239,7 @@ class ZImageTransformer2DModel(CachableDiT, LayerwiseOffloadableModuleMixin):
         Pinning every value a capture consumes keeps those addresses alive;
         contents stay correct because a value never changes for its key.
         Growth is bounded by O(cache sites x captured signatures) small
-        tensors, and nothing is pinned outside graph capture.
+        tensors, and nothing is pinned outside BCG warmup and capture.
         """
         if is_in_breakable_cuda_graph() or (
             _is_cuda and torch.cuda.is_current_stream_capturing()

python/sglang/multimodal_gen/test/unit/test_attention_bcg_outputs.py

@@ -77,14 +77,14 @@ class TestAttentionBCGOutputs(unittest.TestCase):
     def test_local_attention_replays_into_captured_output(self):
         q = torch.zeros(1, 4, 2, 3)
         output = self._capture_and_replay(_attention(layer.LocalAttention), q, q, q)
-        self.assertTrue(torch.all(output == 2))
+        self.assertTrue(torch.all(output == 1))
 
     def test_ulysses_attention_replays_into_captured_tuple(self):
         q = torch.zeros(1, 4, 2, 3)
         replicated_q = torch.zeros(1, 2, 2, 3)
         attention = _attention(layer.UlyssesAttention)
         output, replicated_output = self._capture_and_replay(attention, q, q, q)
-        self.assertTrue(torch.all(output == 2))
+        self.assertTrue(torch.all(output == 1))
         self.assertIsNone(replicated_output)
 
         output, replicated_output = self._capture_and_replay(
@@ -98,8 +98,8 @@ class TestAttentionBCGOutputs(unittest.TestCase):
         )
         self.assertEqual(output.shape, q.shape)
         self.assertEqual(replicated_output.shape, replicated_q.shape)
-        self.assertTrue(torch.all(output == 4))
-        self.assertTrue(torch.all(replicated_output == 4))
+        self.assertTrue(torch.all(output == 2))
+        self.assertTrue(torch.all(replicated_output == 2))
 
 
 if __name__ == "__main__":

python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py

@@ -275,6 +275,9 @@ class TestBCGTPGraphCapture(CustomTestCase):
         self.assertEqual(
             events,
             [
+                # The warmup forwards take the BCG branches outside the TP context.
+                "enter:bcg_enable",
+                "exit:bcg_enable",
                 "enter:tp",
                 "enter:bcg_enable",
                 "enter:bcg_capture",

python/sglang/multimodal_gen/test/unit/test_qwen_image21.py

@@ -926,8 +926,6 @@ def test_prefix_pad_replay_writes_the_captured_output():
             out = QwenImage21Attention.attend_target(attention, q, kv, kv, prefix_pad)
         finally:
             bcg._current_capture_var.reset(token)
-        assert torch.all(out == 4)
-
         prefix_pad.rows.fill_(1)
         with replay_token_scope():
             assert graph._break_fns[0]() is None

python/sglang/srt/layers/attention/hpc_ops_backend.py

@@ -671,7 +671,7 @@ class HPCOpsAttnBackend(AttentionBackend):
         """Graph-splitting wrapper for the fused QKNorm+RoPE+FP8-quant+StoreKV op.
 
         Like ``RadixAttention._eager_attention``, this runs eagerly between captured
-        prefill-graph segments (at capture and at every replay), so the Python
+        prefill-graph segments (at every replay), so the Python
         hand-off of the dynamic Q scales to the following attention op stays
         alive. ``out_q`` is preallocated by the captured segment and mutated in
         place, which is what stitches the surrounding graph segments together.

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

@@ -18,8 +18,8 @@ Each segment is a real torch.cuda.CUDAGraph. Its destructor calls
 releasePool on the shared mempool, so the pool's use_count tracks how
 many segments are alive; the pool stays pinned as long as any segment graph
 is alive. This lets weak_ref_tensor views of intermediate pool-allocated
-tensors remain valid across replays — we don't need Python-managed bridge
-buffers to keep break-point tensors at stable addresses.
+tensors remain valid across replays, so a break's output allocated in the
+captured region needs no other owner.
 """
 
 import functools
@@ -181,46 +181,6 @@ def _weak_ref_if_tensor(x):
     return x
 
 
-def _copy_output(dst: Any, src: Any) -> Any:
-    """Copy src output into dst in-place where possible.
-
-    Handles plain tensors, tuples/lists of tensors, dataclass/object with
-    tensor attributes, and dicts of tensors. Returns dst if in-place copy
-    succeeded, otherwise returns src.
-    """
-    if torch.is_tensor(dst) and torch.is_tensor(src):
-        dst.copy_(src)
-        return dst
-
-    if (
-        isinstance(dst, (tuple, list))
-        and isinstance(src, (tuple, list))
-        and len(dst) == len(src)
-    ):
-        copied = [_copy_output(d, s) for d, s in zip(dst, src)]
-        return tuple(copied) if isinstance(dst, tuple) else copied
-
-    if hasattr(dst, "__dict__") and hasattr(src, "__dict__"):
-        for key, src_val in src.__dict__.items():
-            dst_val = getattr(dst, key, None)
-            if torch.is_tensor(dst_val) and torch.is_tensor(src_val):
-                dst_val.copy_(src_val)
-            else:
-                setattr(dst, key, src_val)
-        return dst
-
-    if isinstance(dst, dict) and isinstance(src, dict):
-        for key, src_val in src.items():
-            dst_val = dst.get(key)
-            if torch.is_tensor(dst_val) and torch.is_tensor(src_val):
-                dst_val.copy_(src_val)
-            else:
-                dst[key] = src_val
-        return dst
-
-    return src
-
-
 class _EagerOnGraphCallable:
     """Bind both graph-aware and explicit eager entry points to the same instance."""
 
@@ -252,9 +212,11 @@ def eager_on_graph(
 ):
     """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.
+    The capture pass records the call without running it; ``capture_stub``, if
+    given, runs at that point for host setup a later captured segment reads.
+    Tensor arguments keep their capture-time identity and are held weakly, so a
+    buffer allocated outside the captured region needs an owner that outlives
+    the graph. Request-specific state is read inside the body when it runs.
 
     Use ``fn.run_eager(...)`` to call the body inside an existing eager region.
     This explicit entry bypasses graph segmentation and capture stubs; methods
@@ -268,52 +230,41 @@ def eager_on_graph(
         return lambda inner: inner
 
     def decorator(inner: Callable):
+        def run(*args, **kwargs) -> None:
+            out = inner(*args, **kwargs)
+            assert out is None, (
+                f"{inner.__qualname__} is a graph break and must return None; "
+                "write the results into buffers the caller allocated"
+            )
+
         @functools.wraps(inner)
         def wrapper(*args, **kwargs):
             capture = _current_capture_var.get()
             if capture is None:
-                return inner(*args, **kwargs)
+                run(*args, **kwargs)
+                return
 
             # End the segment that captured up to this break point.
             capture._end_current_segment()
 
-            # Re-sync ranks after segment teardown (the slow, variable
-            # step) before break fns with rank-coupled collectives and hard
-            # timeouts (DeepEP NORMAL: 100s). Capture-only; replay bypasses
-            # this wrapper.
+            # Re-sync ranks after the slow, variable segment teardown.
             if capture._barrier_fn is not None:
                 capture._barrier_fn()
 
-            # Run the break once so its outputs are allocated and their
-            # addresses recorded. A capture_stub replaces the body during
-            # capture (contents are never consumed; warmup and replay run
-            # the real inner), letting rank-coupled bodies skip the work.
             if capture_stub is not None:
-                output = capture_stub(*args, **kwargs)
-            else:
-                output = inner(*args, **kwargs)
+                capture_stub(*args, **kwargs)
 
             # Weak-ref captured inputs produced by graph segments. Their storage
             # is pinned by the segment CUDAGraphs' mempool use-count, so Python
             # refs do not need to keep every intermediate alive.
-            captured_inner = inner
             captured_args = tuple(_weak_ref_if_tensor(a) for a in args)
             captured_kwargs = {k: _weak_ref_if_tensor(v) for k, v in kwargs.items()}
-            # The eager break output is different: it is allocated between graph
-            # captures and is the static input address consumed by the next
-            # captured segment. Keep a strong reference so replay can safely
-            # copy fresh eager output into that bridge buffer.
-            captured_output = output
-
-            def replay_fn():
-                new_out = captured_inner(*captured_args, **captured_kwargs)
-                return _copy_output(captured_output, new_out)
-
-            capture.cuda_graph._break_fns.append(replay_fn)
+            capture.cuda_graph._break_fns.append(
+                functools.partial(run, *captured_args, **captured_kwargs)
+            )
 
             # Start a fresh CUDAGraph segment for the remainder of the forward.
             capture._begin_new_segment()
-            return output
 
         return functools.update_wrapper(_EagerOnGraphCallable(wrapper, inner), inner)
 

python/sglang/srt/models/deepseek_v4.py

@@ -910,17 +910,7 @@ def deepseek_v4_low_ratio_sources(layer, x, q_lora, positions) -> None:
     )
 
 
-def _low_ratio_sources_capture_stub(layer, x, q_lora, positions) -> None:
-    # A decoder tail graph's capture pass hands the break activations its
-    # segments never computed; the indexer would select candidate blocks from
-    # them and index pages past the live context. Its replay runs the body.
-    if not in_replay_tail_graph():
-        deepseek_v4_low_ratio_sources(layer, x, q_lora, positions)
-
-
-bcg_deepseek_v4_low_ratio_sources = eager_on_graph(
-    deepseek_v4_low_ratio_sources, capture_stub=_low_ratio_sources_capture_stub
-)
+bcg_deepseek_v4_low_ratio_sources = eager_on_graph(deepseek_v4_low_ratio_sources)
 
 
 def _low_ratio_sources_break(attn_backend, forward_batch) -> bool:
@@ -952,8 +942,8 @@ def _write_tail_source_kv_capture_stub(layer, x, positions) -> None:
     get_attn_backend().tail_source_kv_layer_id = layer.layer_id
 
 
-# The capture pass sets only the layer id by which the tail's source layer skips
-# its compressor; the KV the break writes would come from uncomputed inputs.
+# The captured tail reads this layer id to skip the source layer's compressor,
+# so the id is also set while the graph captures.
 bcg_deepseek_v4_write_tail_source_kv = eager_on_graph(
     deepseek_v4_write_tail_source_kv,
     capture_stub=_write_tail_source_kv_capture_stub,
@@ -978,26 +968,27 @@ def _engram_image_token_id(config) -> Optional[int]:
     return None
 
 
-def deepseek_v4_enter_late_layer_tail(sentinel: torch.Tensor) -> torch.Tensor:
+def deepseek_v4_enter_late_layer_tail() -> None:
     # Decoder SWA bounded replay under the prefill CUDA graph: swap the attention
-    # backend onto the tail metadata of the *live* batch (the captured args hold
-    # the capture-time batch). The sentinel keeps the break in the replay sequence.
+    # backend onto the tail metadata of the *live* batch.
     forward_batch = get_tc_piecewise_forward_context().forward_batch
     get_attn_backend().prefill_graph_enter_late_layer_tail(forward_batch)
-    return sentinel
 
 
-def deepseek_v4_exit_late_layer_tail(sentinel: torch.Tensor) -> torch.Tensor:
+def deepseek_v4_exit_late_layer_tail() -> None:
     forward_batch = get_tc_piecewise_forward_context().forward_batch
     get_attn_backend().prefill_graph_exit_late_layer_tail(forward_batch)
-    return sentinel
 
 
-bcg_deepseek_v4_enter_late_layer_tail = eager_on_graph(True)(
-    deepseek_v4_enter_late_layer_tail
+# The captured tail segments read the metadata these swap in, so they also run
+# while the graph captures.
+bcg_deepseek_v4_enter_late_layer_tail = eager_on_graph(
+    deepseek_v4_enter_late_layer_tail,
+    capture_stub=deepseek_v4_enter_late_layer_tail,
 )
-bcg_deepseek_v4_exit_late_layer_tail = eager_on_graph(True)(
-    deepseek_v4_exit_late_layer_tail
+bcg_deepseek_v4_exit_late_layer_tail = eager_on_graph(
+    deepseek_v4_exit_late_layer_tail,
+    capture_stub=deepseek_v4_exit_late_layer_tail,
 )
 
 
@@ -4050,7 +4041,6 @@ class DeepseekV4Model(nn.Module):
                 tail_in_graph=tail_in_graph,
                 late_dp_counts=late_dp_counts,
             )
-        tail_sentinel = positions.new_zeros((1,)) if tail_in_graph else None
         saved_full = None
         # The HIP defer_post gate excludes residual readers. CUDA materializes
         # its producer update explicitly before these readers below.
@@ -4107,7 +4097,7 @@ class DeepseekV4Model(nn.Module):
                     # Eager break: re-run at replay so the backend swaps onto the
                     # live batch's tail metadata; the gathers below are captured
                     # against the persistent index buffers.
-                    tail_sentinel = bcg_deepseek_v4_enter_late_layer_tail(tail_sentinel)
+                    bcg_deepseek_v4_enter_late_layer_tail()
                     saved_full = attn_backend.prefill_graph_late_layer_saved
                     if (
                         self.decoder_replay_tail is not None
@@ -4133,7 +4123,7 @@ class DeepseekV4Model(nn.Module):
                         if capture_dspark:
                             dspark_aux_hidden_states.extend(aux)
                         residual_batch.clear(forward_batch)
-                        bcg_deepseek_v4_exit_late_layer_tail(tail_sentinel)
+                        bcg_deepseek_v4_exit_late_layer_tail()
                         return residual, pre, tail
                 else:
                     saved_full = attn_backend.enter_late_layer_tail(forward_batch)
@@ -4253,7 +4243,7 @@ class DeepseekV4Model(nn.Module):
             exit_late_layer_dp(forward_batch, saved_dp)
         if saved_full is not None:
             if tail_in_graph:
-                tail_sentinel = bcg_deepseek_v4_exit_late_layer_tail(tail_sentinel)
+                bcg_deepseek_v4_exit_late_layer_tail()
             else:
                 attn_backend.exit_late_layer_tail(saved_full, forward_batch)
             return state.residual, state.pre, tail

python/sglang/srt/models/deepseek_v4_replay_tail.py

@@ -665,7 +665,7 @@ def _enter_eager_tail_in_prefill_graph(
     backend, graph_batch: ForwardBatch, source_kv_layer_id: Optional[int]
 ) -> Tuple[tuple, ForwardBatch]:
     """Switches onto exact-row tail metadata; returns the exit state and the
-    batch the tail runs on (the capture batch while the graph captures)."""
+    batch the tail runs on (the dummy batch during the graph warmup)."""
     forward_batch = _unpadded_batch(backend.prefill_graph_replay_batch or graph_batch)
     held = backend.tail_forward_metadata
     backend.tail_forward_metadata = backend._build_late_layer_tail_metadata(
@@ -713,8 +713,8 @@ def _tail_layers_in_prefill_graph(
         num_rows = tail.positions.shape[0]
         dp_rows = None
         if dp:
-            # Exchanged before the forward, except while the prefill graphs capture
-            # (every rank together), when it gathers here.
+            # Exchanged before the forward, except on the dummy batches the prefill
+            # graphs warm up on (every rank together), when it gathers here.
             dp_rows = forward_batch.global_late_layer_num_tokens_cpu
             if dp_rows is None:
                 counts = torch.empty(get_parallel().tp_size, dtype=torch.int64)

test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py

@@ -1,12 +1,3 @@
-"""Tests for the breakable CUDA graph (BCG) runner.
-
-Two test classes:
-- TestBreakableCUDAGraphBasic / TestCopyOutput / TestBreakGraphHelper:
-  unit tests for the core capture / replay mechanism (simple tensor ops).
-- TestBreakableCudaGraph: integration test — spin up Qwen3-8B with
-  --cuda-graph-backend-prefill=breakable and check mgsm_en accuracy.
-"""
-
 import unittest
 from unittest.mock import patch
 
@@ -102,14 +93,15 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         y = torch.zeros(4, device=self.device)
 
         @self.eager_on_graph(enable=True)
-        def eager_op(src):
-            return src * 2.0
+        def eager_op(src, out):
+            out.copy_(src * 2.0)
 
         graph = self.BreakableCUDAGraph()
         stream = torch.cuda.Stream(self.device)
         with self.BreakableCUDAGraphCapture(graph, stream=stream):
             intermediate.copy_(x + 1.0)
-            broken = eager_op(intermediate)
+            broken = torch.empty_like(intermediate)
+            eager_op(intermediate, broken)
             y.copy_(broken + 3.0)
 
         # Replay with new input
@@ -125,20 +117,22 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         y = torch.zeros(4, device=self.device)
 
         @self.eager_on_graph(enable=True)
-        def add_one(src):
-            return src + 1.0
+        def add_one(src, out):
+            out.copy_(src + 1.0)
 
         @self.eager_on_graph(enable=True)
-        def double(src):
-            return src * 2.0
+        def double(src, out):
+            out.copy_(src * 2.0)
 
         graph = self.BreakableCUDAGraph()
         stream = torch.cuda.Stream(self.device)
         with self.BreakableCUDAGraphCapture(graph, stream=stream):
             t1 = x + 1.0  # graph segment 1
-            t2 = add_one(t1)  # break 1: eager
+            t2 = torch.empty_like(t1)
+            add_one(t1, t2)  # break 1: eager
             t3 = t2 + 1.0  # graph segment 2
-            t4 = double(t3)  # break 2: eager
+            t4 = torch.empty_like(t3)
+            double(t3, t4)  # break 2: eager
             y.copy_(t4)  # graph segment 3
 
         # Replay: x=5 -> +1=6 -> add_one=7 -> +1=8 -> double=16
@@ -165,11 +159,12 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         """@eager_on_graph called outside capture should run the function directly."""
 
         @self.eager_on_graph(enable=True)
-        def my_fn(x):
-            return x + 1.0
+        def my_fn(x, out):
+            out.copy_(x + 1.0)
 
         t = torch.tensor([1.0, 2.0], device=self.device)
-        result = my_fn(t)
+        result = torch.empty_like(t)
+        my_fn(t, result)
         self.assertTrue(
             torch.allclose(result, torch.tensor([2.0, 3.0], device=self.device))
         )
@@ -180,14 +175,15 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         y = torch.zeros(4, device=self.device)
 
         @self.eager_on_graph(enable=True)
-        def scale(src):
-            return src * 3.0
+        def scale(src, out):
+            out.copy_(src * 3.0)
 
         graph = self.BreakableCUDAGraph()
         stream = torch.cuda.Stream(self.device)
         with self.BreakableCUDAGraphCapture(graph, stream=stream):
             t = x + 1.0
-            t2 = scale(t)
+            t2 = torch.empty_like(t)
+            scale(t, t2)
             y.copy_(t2)
 
         # First replay: x=0 -> 0+1=1 -> 1*3=3
@@ -208,8 +204,8 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         stream = torch.cuda.Stream(self.device)
 
         @self.eager_on_graph(enable=True)
-        def identity(src):
-            return src
+        def touch(src):
+            pass
 
         graph = self.BreakableCUDAGraph()
         capture_stream = torch.cuda.Stream(self.device)
@@ -217,7 +213,7 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
             stream.wait_stream(torch.cuda.current_stream())
             with torch.cuda.stream(stream):
                 side_output = x + 1.0
-            identity(x)
+            touch(x)
             torch.cuda.current_stream().wait_stream(stream)
             y.copy_(side_output * 2.0)
 
@@ -226,28 +222,6 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         torch.cuda.synchronize()
         self.assertTrue(torch.allclose(y, torch.full((4,), 12.0, device=self.device)))
 
-    def test_eager_output_is_held_strongly_for_replay_bridge(self):
-        """The replay closure must keep the eager output bridge buffer alive."""
-        x = torch.zeros(4, device=self.device)
-        y = torch.zeros(4, device=self.device)
-
-        @self.eager_on_graph(enable=True)
-        def scale(src):
-            return src * 3.0
-
-        graph = self.BreakableCUDAGraph()
-        stream = torch.cuda.Stream(self.device)
-        with self.BreakableCUDAGraphCapture(graph, stream=stream):
-            t = x + 1.0
-            broken = scale(t)
-            y.copy_(broken)
-
-        replay_closure = graph._break_fns[0].__closure__ or ()
-        self.assertTrue(
-            any(cell.cell_contents is broken for cell in replay_closure),
-            "eager output bridge buffer must be strongly captured",
-        )
-
     def test_attention_narrows_padded_positions(self):
 
         num_tokens = 3
@@ -290,64 +264,6 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         self.assertEqual(forward_batch.out_cache_loc.shape[0], padded_num_tokens)
 
 
-class TestCopyOutput(CustomTestCase):
-    """Test the _copy_output helper for structured output writeback."""
-
-    @classmethod
-    def setUpClass(cls):
-        if not torch.cuda.is_available():
-            raise unittest.SkipTest("CUDA not available")
-
-        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
-            _copy_output,
-        )
-
-        cls._copy_output = staticmethod(_copy_output)
-        cls.device = torch.device("cuda:0")
-
-    def test_tensor_copy(self):
-        dst = torch.zeros(4, device=self.device)
-        src = torch.ones(4, device=self.device) * 5.0
-        result = self._copy_output(dst, src)
-        self.assertIs(result, dst)
-        self.assertTrue(torch.allclose(dst, src))
-
-    def test_dict_copy(self):
-        dst = {
-            "a": torch.zeros(4, device=self.device),
-            "b": torch.zeros(4, device=self.device),
-        }
-        src = {
-            "a": torch.ones(4, device=self.device),
-            "b": torch.ones(4, device=self.device) * 2.0,
-        }
-        result = self._copy_output(dst, src)
-        self.assertIs(result, dst)
-        self.assertTrue(torch.allclose(dst["a"], torch.ones(4, device=self.device)))
-        self.assertTrue(
-            torch.allclose(dst["b"], torch.ones(4, device=self.device) * 2.0)
-        )
-
-    def test_object_copy(self):
-        class FakeOutput:
-            def __init__(self, t, label):
-                self.tensor = t
-                self.label = label
-
-        dst = FakeOutput(torch.zeros(4, device=self.device), "old")
-        src = FakeOutput(torch.ones(4, device=self.device) * 3.0, "new")
-        result = self._copy_output(dst, src)
-        self.assertIs(result, dst)
-        self.assertTrue(
-            torch.allclose(dst.tensor, torch.ones(4, device=self.device) * 3.0)
-        )
-        self.assertEqual(dst.label, "new")
-
-    def test_non_tensor_fallback(self):
-        result = self._copy_output(42, 99)
-        self.assertEqual(result, 99)
-
-
 class TestBreakGraphHelper(CustomTestCase):
     """Test the break_graph() convenience function."""
 

test/registered/unit/layers/test_radix_attention.py

@@ -251,15 +251,15 @@ def test_attention_replay_uses_live_batch_and_static_output_buffers(return_lse,
     index_pointer = index_output.data_ptr() if sparse else None
     with attention_scope(live_batch, backend, raw_num_tokens=1):
         graph._break_fns[0]()
-    assert seen == [captured_batch, live_batch]
+    assert seen == [live_batch]
     assert output.data_ptr() == pointer
-    assert torch.all(output[:1] == 2) and torch.count_nonzero(output[1:]) == 0
+    assert torch.all(output[:1] == 1) and torch.count_nonzero(output[1:]) == 0
     if sparse:
         assert index_output.data_ptr() == index_pointer
-        assert torch.all(index_output[:1] == 4)
+        assert torch.all(index_output[:1] == 3)
         assert torch.count_nonzero(index_output[1:]) == 0
     if return_lse:
-        assert torch.all(lse[:1] == 2) and torch.count_nonzero(lse[1:]) == 0
+        assert torch.all(lse[:1] == 1) and torch.count_nonzero(lse[1:]) == 0
 
 
 def test_prefill_graph_policy_restores_after_exception():

test/registered/unit/model_executor/test_eager_on_graph.py

@@ -42,15 +42,17 @@ def recording_capture():
 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
+        def _eager_scale(self, x, out):
+            out.copy_(x * get_forward_batch().scale)
 
     layer = Layer()
     x = torch.tensor([2.0])
+    output = torch.zeros(1)
     with set_forward_batch(SimpleNamespace(scale=3)):
         with recording_capture() as (graph, events):
-            output = layer._eager_scale(x)
+            assert layer._eager_scale(x, output) is None
     assert events == ["end", "begin"]
+    assert output.item() == 0
     pointer = output.data_ptr()
     # Exercise the public, argument-free replay API between two segments.
     graph._segments = [
@@ -73,18 +75,18 @@ def test_method_capture_stub_and_real_replay():
     calls = []
 
     class Layer:
-        def _capture_stub(self, x):
+        def _capture_stub(self, x, out):
             calls.append("stub")
-            return torch.zeros_like(x)
 
         @bcg.eager_on_graph(capture_stub=_capture_stub)
-        def _eager_scale(self, x):
+        def _eager_scale(self, x, out):
             calls.append("real")
-            return x * get_forward_batch().scale
+            out.copy_(x * get_forward_batch().scale)
 
     layer = Layer()
+    output = torch.zeros(2)
     with recording_capture() as (graph, _):
-        output = layer._eager_scale(torch.ones(2))
+        layer._eager_scale(torch.ones(2), output)
     assert calls == ["stub"]
     with set_forward_batch(SimpleNamespace(scale=7)):
         graph._break_fns[0]()
@@ -94,31 +96,46 @@ def test_method_capture_stub_and_real_replay():
 
 def test_plain_callable_and_eager_passthrough():
     @bcg.eager_on_graph
-    def increment(x):
-        return x + 1
+    def increment(x, out):
+        out.copy_(x + 1)
 
-    assert increment(torch.tensor(2)).item() == 3
+    output = torch.tensor(0)
+    increment(torch.tensor(2), output)
+    assert output.item() == 3
     x = torch.tensor(3)
     with recording_capture() as (graph, _):
-        output = increment(x)
+        increment(x, output)
     x.fill_(9)
     graph._break_fns[0]()
     assert output.item() == 10
 
 
+def test_a_break_returning_a_value_fails():
+    @bcg.eager_on_graph
+    def increment(x):
+        return x + 1
+
+    with pytest.raises(AssertionError, match="must return None"):
+        increment(torch.tensor(2))
+    with recording_capture() as (graph, _):
+        increment(torch.tensor(2))
+    with pytest.raises(AssertionError, match="must return None"):
+        graph._break_fns[0]()
+
+
 def test_replay_does_not_retain_capture_or_serving_batches():
     class Batch:
         scale = 3
 
     @bcg.eager_on_graph
-    def scale(x):
-        return x * get_forward_batch().scale
+    def scale(x, out):
+        out.copy_(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))
+            scale(torch.ones(2), torch.zeros(2))
     del batch
     gc.collect()
     assert capture_ref() is None
@@ -136,9 +153,13 @@ def test_replay_does_not_retain_capture_or_serving_batches():
     "decorator", [bcg.eager_on_graph(True), bcg.eager_on_graph(enable=True)]
 )
 def test_legacy_decorator_replays(decorator):
-    increment = decorator(lambda x: x + 1)
+    @decorator
+    def increment(x, out):
+        out.copy_(x + 1)
+
+    output = torch.tensor(0)
     with recording_capture() as (graph, _):
-        output = increment(torch.tensor(2))
+        increment(torch.tensor(2), output)
     graph._break_fns[0]()
     assert output.item() == 3
 
@@ -184,18 +205,19 @@ def test_run_eager_binds_instance_and_keeps_one_outer_break():
             return x * self.scale + offset
 
         @bcg.eager_on_graph
-        def outer(self, x):
-            return self.inner.run_eager(x, offset=1)
+        def outer(self, x, out):
+            out.copy_(self.inner.run_eager(x, offset=1))
 
     first, second = Layer(2), Layer(5)
+    output = torch.zeros(1)
     with recording_capture() as (graph, events):
-        output = first.outer(torch.ones(1))
+        first.outer(torch.ones(1), output)
     assert events == ["end", "begin"]
     assert len(graph._break_fns) == 1
-    assert output.item() == 3
+    assert calls == [] and output.item() == 0
     assert second.inner.run_eager(torch.ones(1), offset=1).item() == 6
     graph._break_fns[0]()
-    assert calls == [2, 5, 2]
+    assert calls == [5, 2]
     assert output.item() == 3
 
 

test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py

@@ -226,9 +226,9 @@ class TestPrefillCudaGraphRunnerHelpers(CustomTestCase):
                         ngram_eos_token_id=None,
                         holder=holder,
                     )
-                else:
-                    for replay in capture.cuda_graph._break_fns:
-                        replay()
+                    self.assertIsNone(holder.batch)
+                for replay in capture.cuda_graph._break_fns:
+                    replay()
                 self.assertEqual(holder.batch.req_indices.tolist(), expected)
 
     def test_dspark_proxy_width_requires_receiving_stage_and_model_support(self):

test/registered/xpu/test_breakable_cuda_graph.py

@@ -80,14 +80,15 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         y = torch.zeros(4, device=self.device)
 
         @self.eager_on_graph(enable=True)
-        def eager_op(src):
-            return src * 2.0
+        def eager_op(src, out):
+            out.copy_(src * 2.0)
 
         graph = self.BreakableCUDAGraph()
         stream = get_device_module().Stream(self.device)
         with self.BreakableCUDAGraphCapture(graph, stream=stream):
             intermediate.copy_(x + 1.0)
-            broken = eager_op(intermediate)
+            broken = torch.empty_like(intermediate)
+            eager_op(intermediate, broken)
             y.copy_(broken + 3.0)
 
         # Replay with new input
@@ -103,20 +104,22 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         y = torch.zeros(4, device=self.device)
 
         @self.eager_on_graph(enable=True)
-        def add_one(src):
-            return src + 1.0
+        def add_one(src, out):
+            out.copy_(src + 1.0)
 
         @self.eager_on_graph(enable=True)
-        def double(src):
-            return src * 2.0
+        def double(src, out):
+            out.copy_(src * 2.0)
 
         graph = self.BreakableCUDAGraph()
         stream = get_device_module().Stream(self.device)
         with self.BreakableCUDAGraphCapture(graph, stream=stream):
             t1 = x + 1.0  # graph segment 1
-            t2 = add_one(t1)  # break 1: eager
+            t2 = torch.empty_like(t1)
+            add_one(t1, t2)  # break 1: eager
             t3 = t2 + 1.0  # graph segment 2
-            t4 = double(t3)  # break 2: eager
+            t4 = torch.empty_like(t3)
+            double(t3, t4)  # break 2: eager
             y.copy_(t4)  # graph segment 3
 
         # Replay: x=5 -> +1=6 -> add_one=7 -> +1=8 -> double=16
@@ -143,11 +146,12 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         """@eager_on_graph called outside capture should run the function directly."""
 
         @self.eager_on_graph(enable=True)
-        def my_fn(x):
-            return x + 1.0
+        def my_fn(x, out):
+            out.copy_(x + 1.0)
 
         t = torch.tensor([1.0, 2.0], device=self.device)
-        result = my_fn(t)
+        result = torch.empty_like(t)
+        my_fn(t, result)
         self.assertTrue(
             torch.allclose(result, torch.tensor([2.0, 3.0], device=self.device))
         )
@@ -158,14 +162,15 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         y = torch.zeros(4, device=self.device)
 
         @self.eager_on_graph(enable=True)
-        def scale(src):
-            return src * 3.0
+        def scale(src, out):
+            out.copy_(src * 3.0)
 
         graph = self.BreakableCUDAGraph()
         stream = get_device_module().Stream(self.device)
         with self.BreakableCUDAGraphCapture(graph, stream=stream):
             t = x + 1.0
-            t2 = scale(t)
+            t2 = torch.empty_like(t)
+            scale(t, t2)
             y.copy_(t2)
 
         # First replay: x=0 -> 0+1=1 -> 1*3=3
@@ -179,86 +184,6 @@ class TestBreakableCUDAGraphBasic(CustomTestCase):
         get_device_module().synchronize()
         self.assertTrue(torch.allclose(y, torch.full((4,), 33.0, device=self.device)))
 
-    def test_eager_output_is_held_strongly_for_replay_bridge(self):
-        """The replay closure must keep the eager output bridge buffer alive."""
-        x = torch.zeros(4, device=self.device)
-        y = torch.zeros(4, device=self.device)
-
-        @self.eager_on_graph(enable=True)
-        def scale(src):
-            return src * 3.0
-
-        graph = self.BreakableCUDAGraph()
-        stream = get_device_module().Stream(self.device)
-        with self.BreakableCUDAGraphCapture(graph, stream=stream):
-            t = x + 1.0
-            broken = scale(t)
-            y.copy_(broken)
-
-        replay_closure = graph._break_fns[0].__closure__ or ()
-        self.assertTrue(
-            any(cell.cell_contents is broken for cell in replay_closure),
-            "eager output bridge buffer must be strongly captured",
-        )
-
-
-class TestCopyOutput(CustomTestCase):
-    """Test the _copy_output helper for structured output writeback."""
-
-    @classmethod
-    def setUpClass(cls):
-        if not get_device_module().is_available():
-            raise unittest.SkipTest(f"{get_device()} not available")
-
-        from sglang.srt.model_executor.runner_backend_utils.breakable_cuda_graph.breakable_cuda_graph import (
-            _copy_output,
-        )
-
-        cls._copy_output = staticmethod(_copy_output)
-        cls.device = torch.device(f"{get_device()}:0")
-
-    def test_tensor_copy(self):
-        dst = torch.zeros(4, device=self.device)
-        src = torch.ones(4, device=self.device) * 5.0
-        result = self._copy_output(dst, src)
-        self.assertIs(result, dst)
-        self.assertTrue(torch.allclose(dst, src))
-
-    def test_dict_copy(self):
-        dst = {
-            "a": torch.zeros(4, device=self.device),
-            "b": torch.zeros(4, device=self.device),
-        }
-        src = {
-            "a": torch.ones(4, device=self.device),
-            "b": torch.ones(4, device=self.device) * 2.0,
-        }
-        result = self._copy_output(dst, src)
-        self.assertIs(result, dst)
-        self.assertTrue(torch.allclose(dst["a"], torch.ones(4, device=self.device)))
-        self.assertTrue(
-            torch.allclose(dst["b"], torch.ones(4, device=self.device) * 2.0)
-        )
-
-    def test_object_copy(self):
-        class FakeOutput:
-            def __init__(self, t, label):
-                self.tensor = t
-                self.label = label
-
-        dst = FakeOutput(torch.zeros(4, device=self.device), "old")
-        src = FakeOutput(torch.ones(4, device=self.device) * 3.0, "new")
-        result = self._copy_output(dst, src)
-        self.assertIs(result, dst)
-        self.assertTrue(
-            torch.allclose(dst.tensor, torch.ones(4, device=self.device) * 3.0)
-        )
-        self.assertEqual(dst.label, "new")
-
-    def test_non_tensor_fallback(self):
-        result = self._copy_output(42, 99)
-        self.assertEqual(result, 99)
-
 
 class TestBreakGraphHelper(CustomTestCase):
     """Test the break_graph() convenience function."""