tensorplay.cuda.graphs.make_graphed_callables

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_args must contain only Tensors whose requires_grad matches 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.