Graph#

API#

class warp_nn.utils.graph.ScopedCapture(
*,
device: str | Device | None = None,
stream: Stream | None = None,
force_module_load: bool | None = None,
external: bool = False,
enabled: bool = True,
)[source]#

Bases: object

Context manager for capturing CUDA graphs.

Adapted from the Warp’s ScopedCapture implementation to support enabling/disabling the capture.

Parameters:

  • device – Data allocation and computation device. If not specified, the default device will be used.

  • stream – CUDA stream to capture on. If not specified, the default stream will be used.

  • force_module_load – Whether to force loading of all kernels before capture.

  • external – Whether the capture was already started externally.

  • enabled – Whether to enable the capture. If disabled, the capture will be skipped and the graph will be None.

property graph: Graph | None[source]#

Captured graph