diff --git a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp
index b6cd4cc..02d1e71 100644
--- a/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp
+++ b/compiler/src/iree/compiler/Dialect/LinalgExt/IR/AggregatedOpInterfaceImpl.cpp
@@ -79,7 +79,8 @@
 }
 
 static Value truncateFloat(OpBuilder &builder, Location loc, AffineMap inputMap,
-                           AffineMap outputMap, Value value, Value output) {
+                           AffineMap outputMap, Value value, Value output,
+                           bool clampToFPRange) {
   SmallVector<AffineMap> compressedMaps =
       compressUnusedDims(SmallVector<AffineMap>{inputMap, outputMap});
   inputMap = compressedMaps[0];
@@ -94,19 +95,23 @@
         auto srcTy = cast<FloatType>(args[0].getType());
         auto dstTy = cast<FloatType>(args[1].getType());
 
-        double mxDbl =
-            APFloat::getLargest(dstTy.getFloatSemantics(), /*Negative=*/false)
-                .convertToDouble();
+        Value input = args[0];
 
-        // Clamp input to dstTy(usually `fp8`) MAX value to prevent NaNs.
-        // We do not clamp for `-MAX` because this function meant to only be
-        // used by attention's exp2 who's value is always > 0.
-        Value mx = builder.create<arith::ConstantOp>(
-            loc, builder.getFloatAttr(srcTy, mxDbl));
-        Value clamp = b.create<arith::MinimumFOp>(loc, mx, args[0]);
+        if (clampToFPRange) {
+          double mxDbl =
+              APFloat::getLargest(dstTy.getFloatSemantics(), /*Negative=*/false)
+                  .convertToDouble();
+
+          // Clamp input to dstTy(usually `fp8`) MAX value to prevent NaNs.
+          // We do not clamp for `-MAX` because this function meant to only be
+          // used by attention's exp2 who's value is always > 0.
+          Value mx = builder.create<arith::ConstantOp>(
+              loc, builder.getFloatAttr(srcTy, mxDbl));
+          input = b.create<arith::MinimumFOp>(loc, mx, input);
+        }
 
         // Convert scale to the same datatype as input.
-        Value trunc = convertScalarToDtype(b, loc, clamp, dstTy,
+        Value trunc = convertScalarToDtype(b, loc, input, dstTy,
                                            /*isUnsignedCast=*/false);
         b.create<linalg::YieldOp>(loc, trunc);
       });
@@ -370,7 +375,8 @@
 
   s = applyPostQKMatmulElementwise(b, loc, getRegion(), s);
 
-  if (qETy.getIntOrFloatBitWidth() <= 8) {
+  bool lowPrecision = qETy.getIntOrFloatBitWidth() <= 8;
+  if (lowPrecision) {
     // For low bit-depth types we perform post Q @ K scaling. This is to avoid
     // losing numerical precision due to the low dynamic range of fp8 types when
     // pre applying the sclaing.
@@ -432,7 +438,7 @@
   auto pETy = getElementTypeOrSelf(p.getType());
   if (pETy != vETy && isa<FloatType>(vETy)) {
     Value convertP = b.create<tensor::EmptyOp>(loc, sSizes, vETy);
-    p = truncateFloat(b, loc, pMap, pMap, p, convertP);
+    p = truncateFloat(b, loc, pMap, pMap, p, convertP, lowPrecision);
   }
 
   Value newAcc = elementwiseValueInPlace<arith::MulFOp>(b, loc, accMap, normMap,