[linalg] Use query shapes for attention broadcast (#4060)

When broadcasting the mask we need to use the query shapes and not the
keys. This is due to the key for GQA having different batch dimensions
than the expanded output.

Author: rsuderman

lib/Conversion/TorchToTMTensor/TorchToTMTensor.cpp

@@ -1841,7 +1841,7 @@ class ConvertAtenScaledDotProductAttentionOp
       int64_t rank = maskTy.getRank();
       bool needsBroadcast = false;
       for (int i = 0, s = rank - 2; i < s; ++i) {
-        needsBroadcast |= maskTy.getDimSize(i) != keyTy.getDimSize(i);
+        needsBroadcast |= maskTy.getDimSize(i) != queryTy.getDimSize(i);
       }
 
       if (needsBroadcast) {
@@ -1850,16 +1850,17 @@ class ConvertAtenScaledDotProductAttentionOp
 
         SmallVector<AffineExpr> maskExprs;
         for (int i = 0, s = rank - 2; i < s; ++i) {
-          maskShape.push_back(keyTy.getDimSize(i));
+          maskShape.push_back(queryTy.getDimSize(i));
 
-          if (maskTy.getDimSize(i) != keyTy.getDimSize(i)) {
+          if (maskTy.getDimSize(i) != queryTy.getDimSize(i)) {
             maskExprs.push_back(rewriter.getAffineConstantExpr(0));
           } else {
             maskExprs.push_back(rewriter.getAffineDimExpr(i));
           }
 
-          if (keyTy.isDynamicDim(i)) {
-            maskDynDims.push_back(rewriter.create<tensor::DimOp>(loc, key, i));
+          if (queryTy.isDynamicDim(i)) {
+            maskDynDims.push_back(
+                rewriter.create<tensor::DimOp>(loc, query, i));
           }
         }
 

projects/pt1/python/torch_mlir_e2e_test/test_suite/basic.py

@@ -5787,7 +5787,7 @@ def __init__(self):
     )
     def forward(self, query, key, value):
         return torch.ops.aten.scaled_dot_product_attention(
-            query, key, value, enable_gqa=True
+            query, key, value, enable_gqa=True, is_causal=True
         )