tensorplay.backends.cuda.can_use_cudnn_attention

tensorplay.backends.cuda.can_use_cudnn_attention(params: _SDPAParams, debug: bool = False) → bool[source]

Check if cudnn_attention can be utilized in scaled_dot_product_attention.

Parameters:

• params – An instance of SDPAParams containing the tensors for query, key, value, an optional attention mask, dropout rate, and a flag indicating if the attention is causal.

• debug – Whether to logging.warn debug information as to why cuDNN attention could not be run. Defaults to False.

Returns:

True if cuDNN can be used with the given parameters; otherwise, False.

Note

This function is dependent on a CUDA-enabled build of TensorPlay. It will return False in non-CUDA environments.