tensorplay.backends.cuda API

Enables or disables fp16/bf16 reduction in math scaled dot product attention.

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.

Check if efficient_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 efficient_attention could not be run. Defaults to False.

Returns:

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

Check if FlashAttention 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 FlashAttention could not be run. Defaults to False.

Returns:

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