diff --git a/docs/docs/advanced_features/breakable_cuda_graph.mdx b/docs/docs/advanced_features/breakable_cuda_graph.mdx index 5bc0589e8e..52763f984e 100644 --- a/docs/docs/advanced_features/breakable_cuda_graph.mdx +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py b/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py index 56dddd4a0e..0c5a219da2 100644 --- a/python/sglang/multimodal_gen/runtime/breakable_cuda_graph/runner.py +++ b/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() diff --git a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py b/python/sglang/multimodal_gen/runtime/layers/attention/layer.py index 82973460a7..b49204ee45 100644 --- a/python/sglang/multimodal_gen/runtime/layers/attention/layer.py +++ b/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 diff --git a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py b/python/sglang/multimodal_gen/runtime/models/dits/zimage.py index 9ffa5a20b3..3d0de2e20b 100644 --- a/python/sglang/multimodal_gen/runtime/models/dits/zimage.py +++ b/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() diff --git a/python/sglang/multimodal_gen/test/unit/test_attention_bcg_outputs.py b/python/sglang/multimodal_gen/test/unit/test_attention_bcg_outputs.py index 814b5d79f6..27b889e16d 100644 --- a/python/sglang/multimodal_gen/test/unit/test_attention_bcg_outputs.py +++ b/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__": diff --git a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py b/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py index ebc591f024..16e61bbe07 100644 --- a/python/sglang/multimodal_gen/test/unit/test_diffusion_bcg_tp_graph_capture.py +++ b/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", diff --git a/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py b/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py index b118720fb6..46e5b74413 100644 --- a/python/sglang/multimodal_gen/test/unit/test_qwen_image21.py +++ b/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 diff --git a/python/sglang/srt/layers/attention/hpc_ops_backend.py b/python/sglang/srt/layers/attention/hpc_ops_backend.py index af37169c88..8ab600f408 100644 --- a/python/sglang/srt/layers/attention/hpc_ops_backend.py +++ b/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. diff --git a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py b/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py index 31fa97ee96..a2e2b34acd 100644 --- a/python/sglang/srt/model_executor/runner_backend_utils/breakable_cuda_graph/breakable_cuda_graph.py +++ b/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) diff --git a/python/sglang/srt/models/deepseek_v4.py b/python/sglang/srt/models/deepseek_v4.py index 19fd546912..8907f4e240 100644 --- a/python/sglang/srt/models/deepseek_v4.py +++ b/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 diff --git a/python/sglang/srt/models/deepseek_v4_replay_tail.py b/python/sglang/srt/models/deepseek_v4_replay_tail.py index 1131e9d853..0ec376e584 100644 --- a/python/sglang/srt/models/deepseek_v4_replay_tail.py +++ b/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) diff --git a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py b/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py index 954a1ecf04..919dc6936b 100644 --- a/test/registered/cuda_graph/breakable/test_breakable_cuda_graph.py +++ b/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.""" diff --git a/test/registered/unit/layers/test_radix_attention.py b/test/registered/unit/layers/test_radix_attention.py index e0b757eb11..c5fc355330 100644 --- a/test/registered/unit/layers/test_radix_attention.py +++ b/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(): diff --git a/test/registered/unit/model_executor/test_eager_on_graph.py b/test/registered/unit/model_executor/test_eager_on_graph.py index 3647ef0db0..eb375d7790 100644 --- a/test/registered/unit/model_executor/test_eager_on_graph.py +++ b/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 diff --git a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py b/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py index 236e646744..8a8c7c3c4b 100644 --- a/test/registered/unit/model_executor/test_prefill_cuda_graph_runner_helpers.py +++ b/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): diff --git a/test/registered/xpu/test_breakable_cuda_graph.py b/test/registered/xpu/test_breakable_cuda_graph.py index 680e2a2595..28453cb6e0 100644 --- a/test/registered/xpu/test_breakable_cuda_graph.py +++ b/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."""