ort::ep::cuda

Struct CUDA

pub struct CUDA { /* private fields */ }

Available on crate feature cuda only.

Expand description

CUDA execution provider for NVIDIA CUDA-enabled GPUs.

Implementations§

Source §

pub fn with_device_id(self, device_id: i32) -> Self

Configures which device the EP should use.

let ep = ep::CUDA::default().with_device_id(0).build();

Source

pub fn with_memory_limit(self, limit: usize) -> Self

Configure the size limit of the device memory arena in bytes.

This only controls how much memory can be allocated to the arena - actual memory usage may be higher due to internal CUDA allocations, like those required for different ConvAlgorithmSearch options.

let ep = ep::CUDA::default().with_memory_limit(2 * 1024 * 1024 * 1024).build();

Source

pub fn with_arena_extend_strategy(self, strategy: ArenaExtendStrategy) -> Self

Configure the strategy for extending the device’s memory arena.

let ep = ep::CUDA::default()
	.with_arena_extend_strategy(ArenaExtendStrategy::SameAsRequested)
	.build();

Source

pub fn with_conv_algorithm_search(self, search: ConvAlgorithmSearch) -> Self

Controls the search mode used to select a kernel for Conv nodes.

cuDNN, the library used by ONNX Runtime’s CUDA EP for many operations, provides many different implementations of the Conv node. Each of these implementations has different performance characteristics depending on the exact hardware and model/input size used. This option controls how cuDNN should determine which implementation to use.

The default search algorithm, Exhaustive, will benchmark all available implementations and use the most performant one. This option is very resource intensive (both computationally on first run and peak-memory-wise), but ensures best performance. It is roughly equivalent to setting torch.backends.cudnn.benchmark = True with PyTorch. See also CUDA::with_conv_max_workspace to configure how much memory the exhaustive search can use (the default is unlimited).

A less resource-intensive option is Heuristic. Rather than benchmarking every implementation, an optimal implementation is chosen based on a set of heuristics, thus saving compute. Heuristic should generally choose an optimal convolution algorithm, except in some corner cases.

Default can also be passed to instruct cuDNN to always use the default implementation (which is rarely the most optimal). Note that the “Default” here refers to the default convolution algorithm being used, it is not the default behavior (that would be Exhaustive).

let ep = ep::CUDA::default()
	.with_conv_algorithm_search(ep::cuda::ConvAlgorithmSearch::Heuristic)
	.build();

Source

pub fn with_conv_max_workspace(self, enable: bool) -> Self

Configure whether the Exhaustive search can use as much memory as it needs.

The default is true. When false, the memory used for the search is limited to 32 MB, which will impact its ability to find an optimal convolution algorithm.

let ep = ep::CUDA::default().with_conv_max_workspace(false).build();

Source

pub fn with_conv1d_pad_to_nc1d(self, enable: bool) -> Self

Configure whether or not to pad 3-dimensional convolutions to [N, C, 1, D] (as opposed to the default [N, C, D, 1]).

Enabling this option might significantly improve performance on devices like the A100. This does not affect convolution operations that do not use 3-dimensional input shapes, or the result of such operations.

let ep = ep::CUDA::default().with_conv1d_pad_to_nc1d(true).build();

Source

pub fn with_cuda_graph(self, enable: bool) -> Self

Configures whether to create a CUDA graph.

CUDA graphs eliminate the overhead of launching kernels sequentially by capturing the launch sequence into a graph that is ‘replayed’ across runs, reducing CPU overhead and possibly improving performance.

Using CUDA graphs comes with limitations, notably:

Models with control flow operators (like If, Loop, or Scan) are not supported.
Input/output shapes cannot change across inference calls.
The address of inputs/outputs cannot change across inference calls, so IoBinding must be used.
Sessions using CUDA graphs are technically not Send or Sync.

Consult the ONNX Runtime documentation on CUDA graphs for more information.

let ep = ep::CUDA::default().with_cuda_graph(true).build();

Source

pub fn with_skip_layer_norm_strict_mode(self, enable: bool) -> Self

Enable ‘strict’ mode for SkipLayerNorm nodes (created via fusion of Add & LayerNorm nodes).

SkipLayerNorm’s strict mode trades performance for accuracy. The default is false (strict mode disabled).

let ep = ep::CUDA::default().with_skip_layer_norm_strict_mode(true).build();

Source

pub fn with_tf32(self, enable: bool) -> Self

Enable the usage of the reduced-precision TensorFloat-32 format for matrix multiplications & convolutions.

TensorFloat-32 is a reduced-precision floating point format available on NVIDIA GPUs since the Ampere microarchitecture. It allows MatMul & Conv to run much faster on Ampere’s Tensor cores. This option is disabled by default.

This option is roughly equivalent to torch.backends.cudnn.allow_tf32 = True & torch.backends.cuda.matmul.allow_tf32 = True or torch.set_float32_matmul_precision("medium") in PyTorch.

let ep = ep::CUDA::default().with_tf32(true).build();

Source

pub fn with_prefer_nhwc(self, enable: bool) -> Self

Configure whether to prefer [N, H, W, C] layout operations over the default [N, C, H, W] layout.

Tensor cores usually operate more efficiently with the NHWC layout, so enabling this option for convolution-heavy models on Tensor core-enabled GPUs may provide a significant performance improvement.

let ep = ep::CUDA::default().with_prefer_nhwc(true).build();

Source

pub unsafe fn with_compute_stream(self, stream: *mut ()) -> Self

Use a custom CUDA device stream rather than the default one.

§Safety

The provided stream must outlive the environment/session configured to use this execution provider.

Source

pub fn with_attention_backend(self, flags: AttentionBackend) -> Self

Configures the available backends used for Attention nodes.

let ep = ep::CUDA::default()
	.with_attention_backend(
		ep::cuda::AttentionBackend::FLASH_ATTENTION | ep::cuda::AttentionBackend::TRT_FUSED_ATTENTION
	)
	.build();

Source