Thank you for adding this! With this PR, can we now pass a 2D attention mask of shape (batch_size, total_sequence_length) directly to MultiHeadAttention as the key_padding_mask input instead of reformatting the attention mask from 2D to 4D causal? If so, can we use the key_padding_mask input this way for the CPU, CUDA, and DML EPs?

For context, the reformatting from 2D to 4D causal was done in the model builder because not all EPs implement causal attention masking with the key_padding_mask input. This caused discrepancies in the outputs from the ONNX model depending on the EP used.

Thank you for adding this! With this PR, can we now pass a 2D attention mask of shape (batch_size, total_sequence_length) directly to MultiHeadAttention as the key_padding_mask input instead of reformatting the attention mask from 2D to 4D causal? If so, can we use the key_padding_mask input this way for the CPU, CUDA, and DML EPs?

For context, the reformatting from 2D to 4D causal was done in the model builder because not all EPs implement causal attention masking with the key_padding_mask input. This caused discrepancies in the outputs from the ONNX model depending on the EP used.

Yes, you can use 2D mask directly (and no need to convert to 4D attention bias). However, 2D mask is for unfused kernel for now. In the future, we might add 2D to 1D conversion to help huggingface model, however that is not best choice for performance since such conversion need extra cuda kernels.

I would suggest to use 1D mask of shape [B] (total seq lengths of each batch assuming right side padding, i.e. the reduce sum of 2D mask) if you want to get benefit of flash attention.

Currently, memory efficient attention need another 1D format with shape [3B+2]. In the future, we will add conversion from 1d mask of shape [B] to tensors compatible with memory efficient attention.