Functions 5
export_dot
functionFull reference ↗graph_pool_handle
functionFull reference ↗graph
functionFull reference ↗- tensorplay.cuda.graphs.graph(cuda_graph, pool=None, stream=None, capture_error_mode='global')[source]
Context-manager that captures CUDA work into a
tensorplay.cuda.CUDAGraph.- Parameters:
cuda_graph (CUDAGraph) – the graph object to capture into.
pool –
Nonegives the graph its own private memory pool; an id fromgraph_pool_handle()(or another graph / itspool_idproperty) shares that pool so allocations from both captures can recycle each other’s space.stream (Stream, optional) – custom capture stream; defaults to the runtime’s dedicated per-device side stream (the legacy default stream cannot capture).
capture_error_mode (str, optional) – see
CUDAGraph.capture_begin().
is_current_stream_capturing
functionFull reference ↗- tensorplay.cuda.graphs.is_current_stream_capturing()[source]
Return True if CUDA graph capture is underway on the current thread.
make_graphed_callables
functionFull reference ↗- tensorplay.cuda.graphs.make_graphed_callables(callables, sample_args, num_warmup_iters=3, allow_unused_input=False, pool=None, capture_error_mode='global')[source]
Callables that run per-iteration with CUDA graph capture.
forward (and backward, via
tensorplay.autograd.grad()) into CUDA graphs sharing one private memory pool, then wraps them in autograd Functions whose forward/backward are graph replays. Per-iteration host overhead drops to two graph launches.carried over verbatim:
sample_argsmust contain only Tensors whoserequires_gradmatches the live workload; modules may not carry hooks or trainable buffers; arguments must keep their order and shapes.- Parameters:
callables – function or
tensorplay.nn.Module, or a tuple of them in live-workload order.sample_args – matching tuple of argument-tuples of CUDA Tensors.
num_warmup_iters – warmup iterations run on the capture stream before capturing (flushes lazy cuDNN/cuBLAS state).
allow_unused_input – passed through to
tensorplay.autograd.grad().pool – share an existing graph pool instead of allocating one.
Classes 1
CUDAGraph
classFull reference ↗- class tensorplay.cuda.graphs.CUDAGraph[source]
- begin_capture_to_if_node(scalar_pred)[source]
Inside an open capture, gate the following work on an
ifnode.scalar_predmust be a single-element CUDA Bool tensor; at replay time the driver samples it and runs the body captured between this call andend_capture_to_conditional_node()only when true.
- begin_capture_to_while_node(scalar_pred)[source]
Like
begin_capture_to_if_node(), but the body loops while the predicate stays true (driver-level while node).
- capture_begin(pool=None, capture_error_mode='global', stream=None)[source]
Begin capture.
- Parameters:
pool –
Nonecaptures into a fresh private pool; an id fromgraph_pool_handle(), another graph’spool_id, or anotherCUDAGraphshares that pool instead.capture_error_mode (str) –
"global"fails the capture if any unsafe CUDA call happens anywhere in the process;"thread_local"only watches this thread;"relaxed"does not guard against unsafe calls.stream – custom capture stream (a
tensorplay.cuda.Stream). Defaults to the runtime’s dedicated per-device side stream; the legacy default stream cannot participate in capture.
- capture_end()[source]
End capture and compile the executable (paid here, not on first replay).
- debug_dump(path)[source]
Write a DOT rendering of the captured graph to
path.Call
enable_debug_mode()before capturing for a dump that includes full node attributes.
- end_capture_to_conditional_node()[source]
Close the open conditional body; subsequent capture returns to the parent stream.
- instantiate()[source]
No-op once instantiated; kept for late callers.
- property pool_id
Allocator pool id this graph captured against.
- replay(stream=None)[source]
Run the graph: launch the cached executable on the current stream.
- Parameters:
stream (Stream, optional) – launch on this explicit stream instead of querying the current one - shaves a TLS lookup off hot loops pinned to a single stream.
- reset()[source]
Destroy the executable and release the pool reference.
All tensors allocated during the capture must be released first.
- set_conditional_handle_for_current_node(scalar_pred)[source]
Refresh the predicate consumed by the innermost open conditional node (used for nested conditionals).
- stage_and_launch(static_inputs, inputs)[source]
Stage every input onto its static buffer and replay in one call.
- Parameters:
static_inputs – buffers captured by the graph (kept alive by the caller).
inputs – fresh tensors whose contents overwrite the matching static buffer this iteration. Contiguous same-dtype/ same-device pairs take a raw async device-to-device copy; anything else falls back to full copy semantics.
This is the low-overhead bulk entry used by
tensorplay._stax.cudagraphs: one Python-to-native crossing for the whole replay instead of one dispatcher round trip per input.
