Support enable_gqa and only support 4D Q, K, and V

[titaiwangms]

1. Support enable_gqa
2. Align PyTorch setting to unsupport Q, K, and V when they are not 4D: https://github.com/pytorch/pytorch/blob/62843c14bbf694f5722fd6e1075da4792507fe42/torch/onnx/_internal/exporter/_torchlib/ops/nn.py#L131-L133

NOTE: torch.nn.functional.scaled_dot_product_attention actually supports 3D, and even Q-3D with K and V - 4D in op tests.

[titaiwangms]

@justinchuby @xadupre Let me know what you think on whether we should support only 4d QKV, or we should fully support whatever torch sdpa supports. Right now, it seems like QKV can have 3d or 4d or even q 3d and kv 4d in torch sdpa.

[justinchuby]

Depending on the ATen op? Does the nn function do preprocessing on the inputs before sending them to the kernel? We just need to support whatever the kernel supports

[codecov]

Codecov Report

❌ Patch coverage is 14.28571% with 24 lines in your changes missing coverage. Please review.
✅ Project coverage is 70.32%. Comparing base (647b22a) to head (e661531).
⚠️ Report is 6 commits behind head on main.

Files with missing lines	Patch %	Lines
onnxscript/function_libs/torch_lib/ops/nn.py	14.28%	22 Missing and 2 partials ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##             main    #2558      +/-   ##
==========================================
- Coverage   70.34%   70.32%   -0.03%     
==========================================
  Files         218      222       +4     
  Lines       26430    26645     +215     
  Branches     2647     2663      +16     
==========================================
+ Hits        18593    18738     +145     
- Misses       6934     6991      +57     
- Partials      903      916      +13     

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

[justinchuby]

Thanks!

[justinchuby]

Could you also add these few lines

if dropout_p > 0.0:
    attn_weight, _ = op.Dropout(attn_weight, dropout_p)

as a micro optimization? Or we can do that separately

[titaiwangms]

Could you also add these few lines

if dropout_p > 0.0:
    attn_weight, _ = op.Dropout(attn_weight, dropout_p)

as a micro optimization? Or we can do that separately

it's already there:

onnxscript/onnxscript/function_libs/torch_lib/ops/nn.py

Line 2041 in 50d7e87

attn_weight, _ = op.Dropout(attn_weight, dropout_p)