[DML EP] Enable more MHA masks

[PatriceVignola]

Those masks are used for MHA in LLaMA.

[sumitsays]

The implementation lgtm. Thank you for adding this.

The contrib op schema for MultiheadAttention says that it only accepts rank <= 2 tensor for mask as mentioned here:

onnxruntime/onnxruntime/core/graph/contrib_ops/bert_defs.cc

Line 955 in f7341e8

"Key padding mask with shape (batch_size) or (3 * batch_size + 2) or (batch_size, kv_sequence_length)",

At the same time, I also see that there is no validation on the rank of this tensor in MultiHeadAttentionTypeAndShapeInference. But CPU EP does throw exception if mask tensor rank > 2.

So I think we should also update the contrib ops doc and may be create a bug for the CPU EP?

[PatriceVignola]

The implementation lgtm. Thank you for adding this.

The contrib op schema for MultiheadAttention says that it only accepts rank <= 2 tensor for mask as mentioned here:

onnxruntime/onnxruntime/core/graph/contrib_ops/bert_defs.cc

Line 955 in f7341e8

"Key padding mask with shape (batch_size) or (3 * batch_size + 2) or (batch_size, kv_sequence_length)",

At the same time, I also see that there is no validation on the rank of this tensor in MultiHeadAttentionTypeAndShapeInference. But CPU EP does throw exception if mask tensor rank > 2.

So I think we should also update the contrib ops doc and may be create a bug for the CPU EP?

The contrib op definitions are being updated in another branch in parallel (including tests), but it will all line up for 1.16.2.

Note: I don't own the other branch (and it's a giant feature branch with hundreds of modified files), which is why I made this separate PR.

[sumitsays]

The implementation lgtm. Thank you for adding this.
The contrib op schema for MultiheadAttention says that it only accepts rank <= 2 tensor for mask as mentioned here:

onnxruntime/onnxruntime/core/graph/contrib_ops/bert_defs.cc

Line 955 in f7341e8

"Key padding mask with shape (batch_size) or (3 * batch_size + 2) or (batch_size, kv_sequence_length)",

At the same time, I also see that there is no validation on the rank of this tensor in MultiHeadAttentionTypeAndShapeInference. But CPU EP does throw exception if mask tensor rank > 2.
So I think we should also update the contrib ops doc and may be create a bug for the CPU EP?

The contrib op definitions are being updated in another branch in parallel (including tests), but it will all line up for 1.16.2.

Note: I don't own the other branch (and it's a giant feature branch with hundreds of modified files), which is why I made this separate PR.

Understood. Thanks again!