[GML] Update SDPA ReduceOpVariants pattern after rebase (#2)

Signed-off-by: James Bartlett <[email protected]>
gimletlabs · Sep 25, 2024 · 5984034 · 5984034
1 parent 9713a52
commit 5984034
Showing 1 changed file with 3 additions and 0 deletions.
diff --git a/lib/Dialect/Torch/Transforms/ReduceOpVariants.cpp b/lib/Dialect/Torch/Transforms/ReduceOpVariants.cpp
@@ -269,6 +269,9 @@ void TorchMatchSpecializedBackendOp::populateSpecializedConversions(
           llvm::SmallVector<Value> newOperands{
               oldOperands[0], oldOperands[1], oldOperands[2], oldOperands[3],
               oldOperands[4], oldOperands[5], oldOperands[7]};
+          Value enableGQA =
+              rewriter.create<ConstantBoolOp>(op->getLoc(), false);
+          newOperands.push_back(enableGQA);
 
           auto newOp = rewriter.create<Torch::AtenScaledDotProductAttentionOp>(
               op.getLoc(), op->getResultTypes()[0], newOperands,