[Torch] Support explicit scale values from `tm_tensor.attention` (#24566)
With dynamically sized GQA, explicit scale settings show up in
ONNX/Torch models such as Qwen3. Consume scale values from TMTensor op
when provided.
Depends on llvm/torch-mlir#4593.
Signed-off-by: Artem Gindinson <gindinson@roofline.ai>
diff --git a/compiler/plugins/input/Torch/InputConversion/ConvertTMTensorToLinalgExt.cpp b/compiler/plugins/input/Torch/InputConversion/ConvertTMTensorToLinalgExt.cpp
index 85121b5..6dcffbc 100644
--- a/compiler/plugins/input/Torch/InputConversion/ConvertTMTensorToLinalgExt.cpp
+++ b/compiler/plugins/input/Torch/InputConversion/ConvertTMTensorToLinalgExt.cpp
@@ -4,6 +4,7 @@
// See https://llvm.org/LICENSE.txt for license information.
// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+#include <mlir/IR/BuiltinAttributes.h>
#include <cstdint>
#include <numeric>
@@ -162,15 +163,24 @@
// have support for batch dims using more general indexing maps, we should
// change this and rely on more general mechanisms.
- // Compute scale = rsqrt(head_dim) in f32.
- int64_t queryRank = op.getQueryType().getRank();
- Value dimIdx =
- rewriter.createOrFold<tensor::DimOp>(loc, query, queryRank - 1);
- Value dimInt = rewriter.createOrFold<arith::IndexCastOp>(
- loc, rewriter.getI64Type(), dimIdx);
- Value dimFloat = rewriter.createOrFold<arith::SIToFPOp>(
- loc, rewriter.getF32Type(), dimInt);
- Value scale = rewriter.createOrFold<math::RsqrtOp>(loc, dimFloat);
+ Value scale;
+ auto scaleAttr = op.getScaleAttr();
+ if (scaleAttr) {
+ auto targetType = cast<FloatType>(op.getQueryType().getElementType());
+ scale = arith::ConstantOp::create(
+ rewriter, loc, targetType,
+ rewriter.getFloatAttr(targetType, scaleAttr.getValueAsDouble()));
+ } else {
+ // Compute scale = rsqrt(head_dim) in f32.
+ int64_t queryRank = op.getQueryType().getRank();
+ Value dimIdx =
+ rewriter.createOrFold<tensor::DimOp>(loc, query, queryRank - 1);
+ Value dimInt = rewriter.createOrFold<arith::IndexCastOp>(
+ loc, rewriter.getI64Type(), dimIdx);
+ Value dimFloat = rewriter.createOrFold<arith::SIToFPOp>(
+ loc, rewriter.getF32Type(), dimInt);
+ scale = rewriter.createOrFold<math::RsqrtOp>(loc, dimFloat);
+ }
// Add batches to standard attention indexing maps.
SmallVector<AffineMap> indexingMaps =
diff --git a/compiler/plugins/input/Torch/InputConversion/test/attention.mlir b/compiler/plugins/input/Torch/InputConversion/test/attention.mlir
index fd58364..fcb6f64 100644
--- a/compiler/plugins/input/Torch/InputConversion/test/attention.mlir
+++ b/compiler/plugins/input/Torch/InputConversion/test/attention.mlir
@@ -59,7 +59,7 @@
// CHECK-LABEL: func.func @attention(
// CHECK-SAME: %[[ARG0:.*]]: tensor<1x3x4xf32>, %[[ARG1:.*]]: tensor<1x3x4xf32>, %[[ARG2:.*]]: tensor<1x3x4xf32>,
-// CHECK: %[[ARG3:.*]]: tensor<1x3x4xf32>) -> tensor<1x3x4xf32> {
+// CHECK-SAME: %[[ARG3:.*]]: tensor<1x3x4xf32>) -> tensor<1x3x4xf32> {
// CHECK: %[[SCALE:.*]] = arith.constant 5.000000e-01 : f32
// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1x3x4xf32>
// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]]} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<1x3x4xf32>, tensor<1x3x4xf32>, tensor<1x3x4xf32>, f32) outs(%[[EMPTY]] : tensor<1x3x4xf32>) {
@@ -69,6 +69,29 @@
// CHECK: return %[[ATTN]] : tensor<1x3x4xf32>
// -----
+func.func @attention_scaled(%arg0: tensor<1x3x4xf32>, %arg1: tensor<1x3x4xf32>, %arg2: tensor<1x3x4xf32>, %arg3: tensor<1x3x4xf32>) -> (tensor<1x3x4xf32>) {
+ %0 = tm_tensor.attention {scale = 0.125 : f64} ins(%arg0, %arg1, %arg2 : tensor<1x3x4xf32>, tensor<1x3x4xf32>, tensor<1x3x4xf32>) outs(%arg3: tensor<1x3x4xf32>) -> tensor<1x3x4xf32>
+ return %0 : tensor<1x3x4xf32>
+}
+
+// CHECK-DAG: #[[$MAP_Q:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d3)>
+// CHECK-DAG: #[[$MAP_K:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d3)>
+// CHECK-DAG: #[[$MAP_V:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d4, d2)>
+// CHECK-DAG: #[[$MAP_S:.+]] = affine_map<(d0, d1, d2, d3, d4) -> ()>
+// CHECK-DAG: #[[$MAP_O:.+]] = affine_map<(d0, d1, d2, d3, d4) -> (d0, d1, d2)>
+
+// CHECK-LABEL: func.func @attention_scaled(
+// CHECK-SAME: %[[ARG0:.*]]: tensor<1x3x4xf32>, %[[ARG1:.*]]: tensor<1x3x4xf32>, %[[ARG2:.*]]: tensor<1x3x4xf32>,
+// CHECK-SAME: %[[ARG3:.*]]: tensor<1x3x4xf32>) -> tensor<1x3x4xf32> {
+// CHECK: %[[SCALE:.*]] = arith.constant 1.250000e-01 : f32
+// CHECK: %[[EMPTY:.*]] = tensor.empty() : tensor<1x3x4xf32>
+// CHECK: %[[ATTN:.*]] = iree_linalg_ext.attention {indexing_maps = [#[[$MAP_Q]], #[[$MAP_K]], #[[$MAP_V]], #[[$MAP_S]], #[[$MAP_O]]]} ins(%[[ARG0]], %[[ARG1]], %[[ARG2]], %[[SCALE]] : tensor<1x3x4xf32>, tensor<1x3x4xf32>, tensor<1x3x4xf32>, f32) outs(%[[EMPTY]] : tensor<1x3x4xf32>) {
+// CHECK: ^[[BLOCK:.+]](%[[SCORE:.+]]: f32):
+// CHECK: linalg_ext.yield %[[SCORE]]
+// CHECK: } -> tensor<1x3x4xf32>
+// CHECK: return %[[ATTN]] : tensor<1x3x4xf32>
+
+// -----
func.func @attention_dyn(%arg0: tensor<?x?x4xf32>, %arg1: tensor<?x?x4xf32>, %arg2: tensor<?x?x4xf32>, %arg3: tensor<?x?x4xf32>) -> (tensor<?x?x4xf32>) {
%0 = tm_tensor.attention ins(%arg0, %arg1, %arg2 : tensor<?x?x4xf32>, tensor<?x?x4xf32>, tensor<?x?x4xf32>) outs(%arg3: tensor<?x?x4xf32>) -> tensor<?x?x4xf32>
return %0 : tensor<?x?x4xf32>
