From 77a7611849ca0615e7c759cd107ef92657116174 Mon Sep 17 00:00:00 2001
From: Umang Yadav <umayadav@amd.com>
Date: Thu, 13 Aug 2026 17:55:52 +0000
Subject: [PATCH 01/10] Reconcile sliding-window clipping with the KV-cache
 matcher

Clip matching returned bounds while sequence-length recovery independently
skipped min/max operations. That could detach a clip from its underlying
currentSeqLen value or fold negative signed bounds into Rock's unsigned
masking semantics.

Return the unclipped input together with validated bounds, accept constants
on either side of commutative min/max operations, and use the same clip
reconstruction for KV-cache and sliding-window masks. Include the matcher
declarations used for constant recognition.

(cherry picked from commit ba1aab52f6a8903455d5a127869d72284a8fbced)
Co-authored-by: Cursor <cursoragent@cursor.com>
---
 mlir/lib/Conversion/TosaToRock/TosaToRock.cpp | 183 ++++++++++--------
 ...-to-rock-attention-sliding-window-neg.mlir |  76 ++++++++
 2 files changed, 173 insertions(+), 86 deletions(-)

diff --git a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
index 09f3448ae452..e2e8f1a45e77 100644
--- a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
+++ b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
@@ -40,6 +40,7 @@
 #include "mlir/IR/BuiltinTypeInterfaces.h"
 #include "mlir/IR/BuiltinTypes.h"
 #include "mlir/IR/Diagnostics.h"
+#include "mlir/IR/Matchers.h"
 #include "mlir/IR/PatternMatch.h"
 #include "mlir/IR/TypeUtilities.h"
 #include "mlir/IR/Types.h"
@@ -2183,18 +2184,71 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     std::optional<int32_t> clipMax;
   };
 
-  // Helper to try detecting KV-cache pattern.
-  // Also detects an optional clip (min(max(x, lo), hi)) on currentSeqLen.
+  struct ClipResult {
+    Value input;
+    int32_t clipMin;
+    int32_t clipMax;
+  };
+
+  // Detect min(max(input, clipMin), clipMax), allowing the constant to appear
+  // on either side of each commutative operation.
+  FailureOr<ClipResult> tryClipPattern(Value input) const {
+    DenseSet<StringRef> expandAndCollapse{
+        tensor::CollapseShapeOp::getOperationName(),
+        tensor::ExpandShapeOp::getOperationName()};
+
+    auto extractI32Constant = [&](Value value) -> std::optional<int32_t> {
+      auto maybeSkipped = getValueSkipping(value, expandAndCollapse);
+      Value unwrapped = succeeded(maybeSkipped) ? *maybeSkipped : value;
+      DenseElementsAttr attr;
+      if (!matchPattern(unwrapped, m_Constant(&attr)) ||
+          !attr.getElementType().isInteger(32) || !attr.isSplat())
+        return std::nullopt;
+      return attr.getSplatValue<int32_t>();
+    };
+
+    auto maybeMin =
+        getDefiningOpSkipping<tosa::MinimumOp>(input, expandAndCollapse);
+    if (failed(maybeMin))
+      return failure();
+
+    Value maxCandidate;
+    std::optional<int32_t> clipMax = extractI32Constant(maybeMin->getInput2());
+    if (clipMax) {
+      maxCandidate = maybeMin->getInput1();
+    } else {
+      clipMax = extractI32Constant(maybeMin->getInput1());
+      if (!clipMax)
+        return failure();
+      maxCandidate = maybeMin->getInput2();
+    }
+
+    auto maybeMax =
+        getDefiningOpSkipping<tosa::MaximumOp>(maxCandidate, expandAndCollapse);
+    if (failed(maybeMax))
+      return failure();
+
+    Value unclippedInput;
+    std::optional<int32_t> clipMin = extractI32Constant(maybeMax->getInput2());
+    if (clipMin) {
+      unclippedInput = maybeMax->getInput1();
+    } else {
+      clipMin = extractI32Constant(maybeMax->getInput1());
+      if (!clipMin)
+        return failure();
+      unclippedInput = maybeMax->getInput2();
+    }
+
+    return ClipResult{unclippedInput, *clipMin, *clipMax};
+  }
+
+  // Helper to try detecting a KV-cache pattern and an optional clip on its
+  // sequence length.
   FailureOr<KVCacheResult>
   tryKVCachePattern(Value input, const DenseSet<StringRef> &seqLenSkip) const {
     DenseSet<StringRef> expandAndCollapse{
         tensor::CollapseShapeOp::getOperationName(),
         tensor::ExpandShapeOp::getOperationName()};
-    DenseSet<StringRef> expandCollapseMinMax{
-        tensor::CollapseShapeOp::getOperationName(),
-        tensor::ExpandShapeOp::getOperationName(),
-        tosa::MaximumOp::getOperationName(),
-        tosa::MinimumOp::getOperationName()};
     FailureOr<Value> maybeNonOne = mulBroadcast(input);
     if (failed(maybeNonOne))
       return failure();
@@ -2209,19 +2263,31 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
         !llvm::all_of(shape.slice(2), [](int32_t v) { return v == 1; }))
       return failure();
 
-    // Try to detect a clip pattern on currentSeqLen before skipping through
-    // min/max. The clip (min(max(x, lo), hi)) may wrap the block argument
-    // and applies to all masks that use currentSeqLen.
     KVCacheResult result;
-    auto maybeClip = tryClipPattern(maybeNonOne.value());
+    Value seqLenCandidate = *maybeNonOne;
+    // MIGraphX may broadcast currentSeqLen more than once, for example first
+    // across heads and then across the key sequence dimension. Peel all
+    // broadcast-only multiplications before looking for a clip.
+    while (true) {
+      FailureOr<Value> maybeInnerBroadcast = mulBroadcast(seqLenCandidate);
+      if (failed(maybeInnerBroadcast))
+        break;
+      seqLenCandidate = *maybeInnerBroadcast;
+    }
+
+    auto maybeClip = tryClipPattern(seqLenCandidate);
     if (succeeded(maybeClip)) {
+      // Rock lowers currentSeqLen masking with unsigned comparisons. Reject a
+      // negative clip instead of dropping it while tracing to the block arg.
+      if (maybeClip->clipMin < 0 || maybeClip->clipMax < 0)
+        return failure();
+      seqLenCandidate = maybeClip->input;
       result.clipMin = maybeClip->clipMin;
       result.clipMax = maybeClip->clipMax;
     }
 
-    // Skip through expand/collapse/min/max to reach the block argument
     auto maybeCurrentSeqLen =
-        getValueSkipping(maybeNonOne.value(), expandCollapseMinMax);
+        getValueSkipping(seqLenCandidate, expandAndCollapse);
     assert(succeeded(maybeCurrentSeqLen) && "Must have non-reshape op");
     Value currentSeqLen = maybeCurrentSeqLen.value();
 
@@ -2229,74 +2295,10 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     if (!isI32BlockArgument(currentSeqLen, seqLenSkip))
       return failure();
 
-    result.seqLen = currentSeqLen;
+    result.seqLen = *maybeCurrentSeqLen;
     return result;
   }
 
-  // Struct for clip detection result
-  struct ClipBounds {
-    int32_t clipMin;
-    int32_t clipMax;
-  };
-
-  // Helper to detect a clip pattern on a value:
-  //   tosa.minimum(tosa.maximum(x, constLo), constHi)
-  FailureOr<ClipBounds> tryClipPattern(Value input) const {
-    DenseSet<StringRef> expandAndCollapse{
-        tensor::CollapseShapeOp::getOperationName(),
-        tensor::ExpandShapeOp::getOperationName()};
-
-    // Helper to extract a splat i32 constant from a value
-    auto extractI32Constant = [&](Value val) -> std::optional<int32_t> {
-      auto maybeSkipped = getValueSkipping(val, expandAndCollapse);
-      Value v = succeeded(maybeSkipped) ? maybeSkipped.value() : val;
-      DenseElementsAttr attr;
-      if (!matchPattern(v, m_Constant(&attr)))
-        return std::nullopt;
-      if (!attr.getElementType().isInteger(32) || !attr.isSplat())
-        return std::nullopt;
-      return attr.getSplatValue<int32_t>();
-    };
-
-    // Look for tosa.minimum (the outer clip op)
-    auto maybeMin =
-        getDefiningOpSkipping<tosa::MinimumOp>(input, expandAndCollapse);
-    if (failed(maybeMin))
-      return failure();
-    auto minOp = maybeMin.value();
-
-    // One input of minimum is a constant (clipMax), the other is maximum
-    Value maxCandidate;
-    std::optional<int32_t> clipMax;
-    clipMax = extractI32Constant(minOp.getInput2());
-    if (clipMax) {
-      maxCandidate = minOp.getInput1();
-    } else {
-      clipMax = extractI32Constant(minOp.getInput1());
-      if (clipMax)
-        maxCandidate = minOp.getInput2();
-      else
-        return failure();
-    }
-
-    // Look for tosa.maximum (the inner clip op)
-    auto maybeMax =
-        getDefiningOpSkipping<tosa::MaximumOp>(maxCandidate, expandAndCollapse);
-    if (failed(maybeMax))
-      return failure();
-    auto maxOp = maybeMax.value();
-
-    // One input of maximum is a constant (clipMin)
-    std::optional<int32_t> clipMin;
-    clipMin = extractI32Constant(maxOp.getInput2());
-    if (!clipMin)
-      clipMin = extractI32Constant(maxOp.getInput1());
-    if (!clipMin)
-      return failure();
-
-    return ClipBounds{*clipMin, *clipMax};
-  }
-
   // Result of sliding-window pattern detection.
   struct SlidingWindowResult {
     int64_t windowSize;
@@ -2314,11 +2316,6 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     DenseSet<StringRef> expandAndCollapse{
         tensor::CollapseShapeOp::getOperationName(),
         tensor::ExpandShapeOp::getOperationName()};
-    DenseSet<StringRef> expandCollapseMinMax{
-        tensor::CollapseShapeOp::getOperationName(),
-        tensor::ExpandShapeOp::getOperationName(),
-        tosa::MaximumOp::getOperationName(),
-        tosa::MinimumOp::getOperationName()};
 
     // Trace through broadcast multiplication (mul by 1)
     FailureOr<Value> maybeNonOne = mulBroadcast(input);
@@ -2365,8 +2362,21 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
     std::optional<int32_t> clipMin;
     std::optional<int32_t> clipMax;
-    auto maybeClip = tryClipPattern(seqLenOperand);
+    Value seqLenCandidate = seqLenOperand;
+    while (true) {
+      FailureOr<Value> maybeInnerBroadcast = mulBroadcast(seqLenCandidate);
+      if (failed(maybeInnerBroadcast))
+        break;
+      seqLenCandidate = *maybeInnerBroadcast;
+    }
+
+    auto maybeClip = tryClipPattern(seqLenCandidate);
     if (succeeded(maybeClip)) {
+      // Rock lowers currentSeqLen masking with unsigned comparisons. Reject a
+      // negative clip instead of dropping it while tracing to the block arg.
+      if (maybeClip->clipMin < 0 || maybeClip->clipMax < 0)
+        return failure();
+      seqLenCandidate = maybeClip->input;
       clipMin = maybeClip->clipMin;
       clipMax = maybeClip->clipMax;
     }
@@ -2374,8 +2384,9 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     // An unrelated greater(x - const, col) is not a sliding-window mask. The
     // non-constant operand must resolve to an i32 currentSeqLen block argument.
     FailureOr<Value> maybeSeqLen =
-        getValueSkipping(seqLenOperand, expandCollapseMinMax);
-    Value seqLen = succeeded(maybeSeqLen) ? maybeSeqLen.value() : seqLenOperand;
+        getValueSkipping(seqLenCandidate, expandAndCollapse);
+    Value seqLen =
+        succeeded(maybeSeqLen) ? maybeSeqLen.value() : seqLenCandidate;
     if (!isI32BlockArgument(seqLen, seqLenSkip))
       return failure();
 
diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
index f0a8320536d4..e35bea91444a 100644
--- a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
+++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
@@ -208,6 +208,82 @@ func.func @sliding_window_kvcache_mismatched_clip(%arg0: tensor<1xi32>, %arg1: t
 
 // -----
 
+// A negative currentSeqLen clip cannot be represented by Rock's unsigned mask
+// comparisons. Keep both masks explicit rather than dropping the clamp while
+// folding them into rock.attention.
+// CHECK-LABEL: func @sliding_window_kvcache_negative_clip
+// CHECK: rock.attention
+// CHECK-NOT: slidingWindowSize
+// CHECK-NOT: currentSeqLen
+// CHECK: qk = elementwise
+// CHECK: tosa.maximum
+// CHECK: tosa.minimum
+func.func @sliding_window_kvcache_negative_clip(%arg0: tensor<1xi32>, %arg1: tensor<12xf16>, %arg2: tensor<32xf16>, %arg3: tensor<32xf16>) -> tensor<4xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} {
+  %0 = "tosa.const"() <{values = dense<4> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32>
+  %2 = "tosa.const"() <{values = dense<-1> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32>
+  %4 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x2x1x8xf32>}> : () -> tensor<1x2x1x8xf32>
+  %5 = "tosa.const"() <{values = dense<1> : tensor<1x2x1x8xi8>}> : () -> tensor<1x2x1x8xi8>
+  %7 = "tosa.const"() <{values = dense<1> : tensor<8x1x1x1xi32>}> : () -> tensor<8x1x1x1xi32>
+  %8 = "tosa.const"() <{values = dense<5.000000e-01> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16>
+  %9 = "tosa.const"() <{values = dense<0xFC00> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16>
+  %11 = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16>
+  %16 = "tosa.const"() <{values = dense<-3> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32>
+  %17 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %18 = "tosa.const"() <{values = dense<1> : tensor<1x1x1x8xi32>}> : () -> tensor<1x1x1x8xi32>
+  %20 = "tosa.const"() <{values = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xi32>}> : () -> tensor<8xi32>
+  %cst = arith.constant dense<[[[[0, 1, 2, 3, 4, 5, 6, 7]]]]> : tensor<1x1x1x8xi32>
+  %expanded = tensor.expand_shape %arg2 [[0, 1, 2, 3]] output_shape [1, 2, 8, 2] : tensor<32xf16> into tensor<1x2x8x2xf16>
+  %expanded_0 = tensor.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [1, 6, 1, 2] : tensor<12xf16> into tensor<1x6x1x2xf16>
+  %expanded_1 = tensor.expand_shape %arg0 [[0, 1, 2, 3]] output_shape [1, 1, 1, 1] : tensor<1xi32> into tensor<1x1x1x1xi32>
+  %k23 = tosa.maximum %expanded_1, %2 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
+  %k24 = tosa.minimum %k23, %0 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
+  %23 = tosa.maximum %expanded_1, %2 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
+  %24 = tosa.minimum %23, %0 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
+  %extracted_slice = tensor.extract_slice %expanded_0[0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1] : tensor<1x6x1x2xf16> to tensor<1x2x1x2xf16>
+  %25 = tosa.transpose %expanded {perms = array<i32: 0, 1, 3, 2>} : (tensor<1x2x8x2xf16>) -> tensor<1x2x2x8xf16>
+  %collapsed = tensor.collapse_shape %extracted_slice [[0, 1], [2], [3]] : tensor<1x2x1x2xf16> into tensor<2x1x2xf16>
+  %collapsed_2 = tensor.collapse_shape %25 [[0, 1], [2], [3]] : tensor<1x2x2x8xf16> into tensor<2x2x8xf16>
+  %26 = tosa.matmul %collapsed, %collapsed_2, %11, %11 {acc_type = f32} : (tensor<2x1x2xf16>, tensor<2x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x8xf16>
+  %expanded_3 = tensor.expand_shape %26 [[0, 1], [2], [3]] output_shape [1, 2, 1, 8] : tensor<2x1x8xf16> into tensor<1x2x1x8xf16>
+  %27 = tosa.mul %expanded_3, %8, %17 : (tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>, tensor<1xi8>) -> tensor<1x2x1x8xf16>
+  %k36 = tosa.mul %k24, %18, %17 : (tensor<1x1x1x1xi32>, tensor<1x1x1x8xi32>, tensor<1xi8>) -> tensor<1x1x1x8xi32>
+  %k37 = tosa.greater %cst, %k36 : (tensor<1x1x1x8xi32>, tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi1>
+  %k38 = tosa.cast %k37 : (tensor<1x1x1x8xi1>) -> tensor<1x1x1x8xi32>
+  %k39 = tosa.cast %k38 : (tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi8>
+  %k40 = tosa.mul %k39, %5, %17 : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
+  %k41 = tosa.cast %k40 : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
+  %kvsel = tosa.select %k41, %9, %27 : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %28 = tosa.add %24, %16 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
+  %29 = tosa.mul %28, %7, %17 : (tensor<1x1x1x1xi32>, tensor<8x1x1x1xi32>, tensor<1xi8>) -> tensor<8x1x1x1xi32>
+  %collapsed_4 = tensor.collapse_shape %29 [[0, 1, 2, 3]] : tensor<8x1x1x1xi32> into tensor<8xi32>
+  %30 = tosa.greater %collapsed_4, %20 : (tensor<8xi32>, tensor<8xi32>) -> tensor<8xi1>
+  %31 = tosa.cast %30 : (tensor<8xi1>) -> tensor<8xi32>
+  %32 = tosa.cast %31 : (tensor<8xi32>) -> tensor<8xi8>
+  %expanded_5 = tensor.expand_shape %32 [[0, 1, 2, 3]] output_shape [1, 1, 1, 8] : tensor<8xi8> into tensor<1x1x1x8xi8>
+  %33 = tosa.mul %expanded_5, %5, %17 : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
+  %34 = tosa.cast %33 : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
+  %35 = tosa.select %34, %9, %kvsel : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %43 = tosa.cast %35 : (tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf32>
+  %44 = tosa.reduce_max %43 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
+  %45 = tosa.mul %44, %4, %17 : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %46 = tosa.sub %43, %45 : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %47 = tosa.exp %46 : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %48 = tosa.reduce_sum %47 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
+  %49 = tosa.mul %48, %4, %17 : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %50 = tosa.reciprocal %49 : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %51 = tosa.mul %47, %50, %17 : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %52 = tosa.cast %51 : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf16>
+  %collapsed_6 = tensor.collapse_shape %52 [[0, 1], [2], [3]] : tensor<1x2x1x8xf16> into tensor<2x1x8xf16>
+  %expanded_7 = tensor.expand_shape %arg3 [[0, 1, 2]] output_shape [2, 8, 2] : tensor<32xf16> into tensor<2x8x2xf16>
+  %53 = tosa.matmul %collapsed_6, %expanded_7, %11, %11 {acc_type = f32} : (tensor<2x1x8xf16>, tensor<2x8x2xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x2xf16>
+  %expanded_8 = tensor.expand_shape %53 [[0, 1], [2], [3]] output_shape [1, 2, 1, 2] : tensor<2x1x2xf16> into tensor<1x2x1x2xf16>
+  %54 = tosa.transpose %expanded_8 {perms = array<i32: 0, 2, 1, 3>} : (tensor<1x2x1x2xf16>) -> tensor<1x1x2x2xf16>
+  %collapsed_9 = tensor.collapse_shape %54 [[0, 1, 2, 3]] : tensor<1x1x2x2xf16> into tensor<4xf16>
+  return %collapsed_9 : tensor<4xf16>
+}
+
+// -----
+
 // A greater(x - window, col) mask whose x is not currentSeqLen must not be
 // classified as sliding-window attention.
 // CHECK-LABEL: func @not_sliding_window_wrong_operand

From bc36bb0c1a9dcdebd3b4c4b10fb42ab469ff79f6 Mon Sep 17 00:00:00 2001
From: Umang Yadav <umayadav@amd.com>
Date: Thu, 13 Aug 2026 17:55:53 +0000
Subject: [PATCH 02/10] Fix sliding-window sequence-length reconstruction

MIGraphX may broadcast currentSeqLen through multiple multiply-by-one layers.
Peeling only the outer broadcast left an intermediate tensor instead of the
original block argument, so matching could fail or rebuild the wrong head
shape.

Peel every broadcast-only multiplication before resolving currentSeqLen, then
reconstruct the required head broadcast from the validated block argument.

(cherry picked from commit 3642fef9e4dc9f8a69e980047714bdb713e59e48)
Co-authored-by: Cursor <cursoragent@cursor.com>
---
 mlir/lib/Conversion/TosaToRock/TosaToRock.cpp | 50 ++++++++-----------
 .../tosa-to-rock-attention-kvcache.mlir       |  4 +-
 2 files changed, 25 insertions(+), 29 deletions(-)

diff --git a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
index e2e8f1a45e77..c10aca674a39 100644
--- a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
+++ b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
@@ -2190,6 +2190,16 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     int32_t clipMax;
   };
 
+  // Peel all multiply-by-one operations used to broadcast a scalar-like value.
+  Value peelBroadcasts(Value input) const {
+    while (true) {
+      FailureOr<Value> maybeNonOne = mulBroadcast(input);
+      if (failed(maybeNonOne))
+        return input;
+      input = *maybeNonOne;
+    }
+  }
+
   // Detect min(max(input, clipMin), clipMax), allowing the constant to appear
   // on either side of each commutative operation.
   FailureOr<ClipResult> tryClipPattern(Value input) const {
@@ -2246,9 +2256,6 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
   // sequence length.
   FailureOr<KVCacheResult>
   tryKVCachePattern(Value input, const DenseSet<StringRef> &seqLenSkip) const {
-    DenseSet<StringRef> expandAndCollapse{
-        tensor::CollapseShapeOp::getOperationName(),
-        tensor::ExpandShapeOp::getOperationName()};
     FailureOr<Value> maybeNonOne = mulBroadcast(input);
     if (failed(maybeNonOne))
       return failure();
@@ -2268,12 +2275,7 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     // MIGraphX may broadcast currentSeqLen more than once, for example first
     // across heads and then across the key sequence dimension. Peel all
     // broadcast-only multiplications before looking for a clip.
-    while (true) {
-      FailureOr<Value> maybeInnerBroadcast = mulBroadcast(seqLenCandidate);
-      if (failed(maybeInnerBroadcast))
-        break;
-      seqLenCandidate = *maybeInnerBroadcast;
-    }
+    seqLenCandidate = peelBroadcasts(seqLenCandidate);
 
     auto maybeClip = tryClipPattern(seqLenCandidate);
     if (succeeded(maybeClip)) {
@@ -2286,13 +2288,13 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       result.clipMax = maybeClip->clipMax;
     }
 
-    auto maybeCurrentSeqLen =
-        getValueSkipping(seqLenCandidate, expandAndCollapse);
-    assert(succeeded(maybeCurrentSeqLen) && "Must have non-reshape op");
-    Value currentSeqLen = maybeCurrentSeqLen.value();
-
-    // Verify currentSeqLen is i32 and traces back to a block argument
-    if (!isI32BlockArgument(currentSeqLen, seqLenSkip))
+    // Resolve through the remaining reshape, transpose, and broadcast chain.
+    // Returning the block argument lets addBroadcastForBlockArg reconstruct
+    // the head broadcast after nested broadcasts have been peeled.
+    FailureOr<Value> maybeCurrentSeqLen =
+        getValueSkipping(seqLenCandidate, seqLenSkip);
+    if (failed(maybeCurrentSeqLen) ||
+        !isI32BlockArgument(*maybeCurrentSeqLen, seqLenSkip))
       return failure();
 
     result.seqLen = *maybeCurrentSeqLen;
@@ -2362,13 +2364,7 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
     std::optional<int32_t> clipMin;
     std::optional<int32_t> clipMax;
-    Value seqLenCandidate = seqLenOperand;
-    while (true) {
-      FailureOr<Value> maybeInnerBroadcast = mulBroadcast(seqLenCandidate);
-      if (failed(maybeInnerBroadcast))
-        break;
-      seqLenCandidate = *maybeInnerBroadcast;
-    }
+    Value seqLenCandidate = peelBroadcasts(seqLenOperand);
 
     auto maybeClip = tryClipPattern(seqLenCandidate);
     if (succeeded(maybeClip)) {
@@ -2384,13 +2380,11 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     // An unrelated greater(x - const, col) is not a sliding-window mask. The
     // non-constant operand must resolve to an i32 currentSeqLen block argument.
     FailureOr<Value> maybeSeqLen =
-        getValueSkipping(seqLenCandidate, expandAndCollapse);
-    Value seqLen =
-        succeeded(maybeSeqLen) ? maybeSeqLen.value() : seqLenCandidate;
-    if (!isI32BlockArgument(seqLen, seqLenSkip))
+        getValueSkipping(seqLenCandidate, seqLenSkip);
+    if (failed(maybeSeqLen) || !isI32BlockArgument(*maybeSeqLen, seqLenSkip))
       return failure();
 
-    return SlidingWindowResult{maybeWindowSize.value(), seqLen, clipMin,
+    return SlidingWindowResult{maybeWindowSize.value(), *maybeSeqLen, clipMin,
                                clipMax};
   }
 
diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
index 61ec7ebe34b3..a142725c7633 100644
--- a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
+++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
@@ -24,7 +24,9 @@ func.func @mlir_attention(%arg0: tensor<12288xf16> {mhal.read_access}, %arg1: te
   %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8> 
   %7 = tosa.mul %cst, %6, %shift : (tensor<1x1x1x1024xi32>, tensor<1x32x1x1024xi32>, tensor<1xi8>) -> tensor<1x32x1x1024xi32>
   %expanded_5 = tensor.expand_shape %arg3 [[0, 1, 2, 3]] output_shape [1, 32, 1, 1] : tensor<32xi32> into tensor<1x32x1x1xi32>
-  %8 = tosa.mul %expanded_5, %6, %shift : (tensor<1x32x1x1xi32>, tensor<1x32x1x1024xi32>, tensor<1xi8>) -> tensor<1x32x1x1024xi32>
+  %seq_one = "tosa.const"() <{values = dense<1> : tensor<1x32x1x1xi32>}> : () -> tensor<1x32x1x1xi32>
+  %seq_broadcast = tosa.mul %expanded_5, %seq_one, %shift : (tensor<1x32x1x1xi32>, tensor<1x32x1x1xi32>, tensor<1xi8>) -> tensor<1x32x1x1xi32>
+  %8 = tosa.mul %seq_broadcast, %6, %shift : (tensor<1x32x1x1xi32>, tensor<1x32x1x1024xi32>, tensor<1xi8>) -> tensor<1x32x1x1024xi32>
   %9 = tosa.greater %7, %8 : (tensor<1x32x1x1024xi32>, tensor<1x32x1x1024xi32>) -> tensor<1x32x1x1024xi1>
   %10 = tosa.cast %9 : (tensor<1x32x1x1024xi1>) -> tensor<1x32x1x1024xi32>
   %11 = tosa.cast %10 : (tensor<1x32x1x1024xi32>) -> tensor<1x32x1x1024xi8>

From e1f50d91eaa2f63d79550071b2e424e5b87d07d6 Mon Sep 17 00:00:00 2001
From: Umang Yadav <umayadav@amd.com>
Date: Thu, 13 Aug 2026 17:55:54 +0000
Subject: [PATCH 03/10] Preserve sliding-window-only mask semantics

A sliding-window lower bound does not prove that keys after currentSeqLen are
invalid. Folding that mask alone by adopting its operand as currentSeqLen
introduced a KV-cache upper mask that was absent from the source IR.

Require an independently matched KV-cache upper mask before folding the window.
Otherwise retain the lower-bound select in the elementwise region, preserving
the original semantics. Cover this behavior in the negative matcher test.

(cherry picked from commit 105932404a1b7b264640f2664ac9798d8b151828)
Co-authored-by: Cursor <cursoragent@cursor.com>
---
 mlir/lib/Conversion/TosaToRock/TosaToRock.cpp | 58 +++++++++++++------
 ...-to-rock-attention-sliding-window-neg.mlir | 17 ++++--
 2 files changed, 52 insertions(+), 23 deletions(-)

diff --git a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
index c10aca674a39..d76768380dc9 100644
--- a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
+++ b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
@@ -2116,9 +2116,18 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     // Clip bounds detected on currentSeqLen during KV-cache pattern matching.
     std::optional<int32_t> seqLenClipMin;
     std::optional<int32_t> seqLenClipMax;
-    // The currentSeqLen block argument and clip referenced by the
-    // sliding-window mask.
+    // The currentSeqLen block argument referenced by the sliding-window mask.
+    // Used to verify that the sliding-window and KV-cache masks have the same
+    // position operand.
     Value slidingWindowSeqLen;
+    // Clip bounds detected on the sliding-window seq-len operand.
+    //
+    // In valid IR the seq-len is clamped once and that single clip (the same
+    // min/max ops) feeds every mask, so when both a KV-cache and a
+    // sliding-window mask are present these bounds are identical to
+    // seqLenClip{Min,Max}. The two pairs exist only because each mask is
+    // matched independently; they are reconciled (and a mismatch is rejected)
+    // in getSeqLenMask so that only one effective clip is ever emitted.
     std::optional<int32_t> slidingWindowClipMin;
     std::optional<int32_t> slidingWindowClipMax;
   };
@@ -2304,7 +2313,12 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
   // Result of sliding-window pattern detection.
   struct SlidingWindowResult {
     int64_t windowSize;
+    // The currentSeqLen operand feeding (currentSeqLen - windowSize), resolved
+    // through reshape/clip ops so it can be matched against the KV-cache
+    // seq-len.
     Value seqLen;
+    // Clip bounds (min(max(x, lo), hi)) detected on the seq-len operand.
+    // Carried so they can be compared with the KV-cache clip bounds.
     std::optional<int32_t> clipMin;
     std::optional<int32_t> clipMax;
   };
@@ -2362,6 +2376,9 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     if (failed(maybeWindowSize))
       return failure();
 
+    // The seq-len operand may be wrapped in a clip (min(max(x, lo), hi)) just
+    // like the KV-cache path. Detect it before skipping through the min/max so
+    // its bounds can be checked against the KV-cache mask.
     std::optional<int32_t> clipMin;
     std::optional<int32_t> clipMax;
     Value seqLenCandidate = peelBroadcasts(seqLenOperand);
@@ -2588,7 +2605,11 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     if (succeeded(isConstantRange(greater.getInput2(), 0))) {
       Value input1 = greater.getInput1();
 
-      // Try sliding window pattern if not already found
+      // Try sliding window pattern if not already found. Record the operand and
+      // any clip bounds; the consistency check against a KV-cache seq-len is
+      // done once, centrally, after all masks have been peeled. Doing it here
+      // would miss the case where the sliding-window mask is seen before the
+      // KV-cache mask.
       if (!result.slidingWindowSize) {
         auto maybeSlidingWindow = trySlidingWindowPattern(input1, seqLenSkip);
         if (succeeded(maybeSlidingWindow)) {
@@ -2662,20 +2683,23 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     // the validated operand after all masks have been analyzed so the result is
     // independent of the select nesting order.
     if (currentResult.slidingWindowSize) {
-      if (currentResult.seqLen) {
-        if (!sameSeqLenBlockArg(currentResult.seqLen,
-                                currentResult.slidingWindowSeqLen, seqLenSkip))
-          return failure();
-        // A single attention op cannot represent different clamps for the
-        // KV-cache and sliding-window masks.
-        if (currentResult.seqLenClipMin != currentResult.slidingWindowClipMin ||
-            currentResult.seqLenClipMax != currentResult.slidingWindowClipMax)
-          return failure();
-      } else {
-        currentResult.seqLen = currentResult.slidingWindowSeqLen;
-        currentResult.seqLenClipMin = currentResult.slidingWindowClipMin;
-        currentResult.seqLenClipMax = currentResult.slidingWindowClipMax;
-      }
+      // Sliding-window folding requires both the lower window bound and the
+      // KV-cache upper bound represented by the attention op.
+      if (!currentResult.seqLen)
+        return failure();
+
+      // The sliding-window mask must reference the same seq-len block argument
+      // as the KV-cache mask; otherwise the two masks disagree on the sequence
+      // length and cannot be folded into an op with one currentSeqLen.
+      if (!sameSeqLenBlockArg(currentResult.seqLen,
+                              currentResult.slidingWindowSeqLen, seqLenSkip))
+        return failure();
+      // Both masks must clamp currentSeqLen identically. sameSeqLenBlockArg
+      // only matches the underlying block argument (it skips through the clip
+      // min/max), so a divergent clip would otherwise be silently dropped.
+      if (currentResult.seqLenClipMin != currentResult.slidingWindowClipMin ||
+          currentResult.seqLenClipMax != currentResult.slidingWindowClipMax)
+        return failure();
     }
 
     // We need at least one pattern to be detected
diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
index e35bea91444a..85a8953ef1f0 100644
--- a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
+++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
@@ -1,13 +1,18 @@
 // RUN: sed s/##TOKEN_ARCH##/%arch/g %s | rocmlir-opt --tosa-to-rock -split-input-file -verify-diagnostics -o -| FileCheck %s
 
-// A sliding-window mask without a separate KV-cache mask must adopt its
-// validated seq-len operand as currentSeqLen and preserve its clip.
+// Edge case 1: a sliding-window-shaped lower mask WITHOUT a separate KV-cache
+// upper mask. Folding this as sliding-window attention would set currentSeqLen
+// and introduce a new upper mask for keys after that position. Keep the lower
+// mask in the elementwise region instead.
 // CHECK-LABEL: func @sliding_window_no_kvcache
-// CHECK: %[[MAX:.*]] = tosa.maximum
-// CHECK: %[[CLIP:.*]] = tosa.minimum %[[MAX]]
 // CHECK: rock.attention
-// CHECK: currentSeqLen = (%[[CLIP]]
-// CHECK: slidingWindowSize = 3
+// CHECK-NOT: currentSeqLen
+// CHECK-NOT: slidingWindowSize
+// CHECK: qk = elementwise
+// CHECK: tosa.maximum
+// CHECK: tosa.minimum
+// CHECK: tosa.greater
+// CHECK: tosa.select
 func.func @sliding_window_no_kvcache(%arg0: tensor<1xi32>, %arg1: tensor<12xf16>, %arg2: tensor<32xf16>, %arg3: tensor<32xf16>) -> tensor<4xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} {
   %0 = "tosa.const"() <{values = dense<4> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32>
   %4 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x2x1x8xf32>}> : () -> tensor<1x2x1x8xf32>

From 50b2b9f91da114d637dcb327ea2615cb804966c9 Mon Sep 17 00:00:00 2001
From: Umang Yadav <umayadav@amd.com>
Date: Thu, 13 Aug 2026 17:55:55 +0000
Subject: [PATCH 04/10] Default sliding-window tuning to full cache

currentSeqLen is runtime data and is intentionally omitted from attention
tuning keys, but rocmlir-gen required it whenever sliding_window_size was set.
Reconstructing a serialized tuning problem therefore failed before generating
valid attention IR.

When the runtime position is absent, synthesize seq_len_k - 1 for every group.
This represents a full cache without changing tuning identity. Update option,
problem-key, bufferized generation, and performance-runner coverage accordingly.

(cherry picked from commit a57d63d24488d94b9338ff61f4cb171224d02483)
Co-authored-by: Cursor <cursoragent@cursor.com>
---
 .../rocmlir-gen/attention-sliding-window.mlir |  9 ++++++
 mlir/test/rocmlir-gen/options.mlir            |  8 ++---
 mlir/test/rocmlir-gen/problem-key.mlir        |  6 ++--
 mlir/tools/rocmlir-gen/rocmlir-gen.cpp        | 32 +++++++++++++------
 mlir/utils/performance/perfRunner.py          |  3 +-
 5 files changed, 40 insertions(+), 18 deletions(-)

diff --git a/mlir/test/rocmlir-gen/attention-sliding-window.mlir b/mlir/test/rocmlir-gen/attention-sliding-window.mlir
index 27348f354321..4791b0952b68 100644
--- a/mlir/test/rocmlir-gen/attention-sliding-window.mlir
+++ b/mlir/test/rocmlir-gen/attention-sliding-window.mlir
@@ -1,5 +1,6 @@
 // RUN: rocmlir-gen --arch gfx90a:sramecc+:xnack- --operation attention -current_seq_len=33 -sliding_window_size=16 -seq_len_q 1 -seq_len_k 64 -head_dim_qk 32 -head_dim_v 32 -t f32 -pv --apply-bufferization-pipeline=false | rocmlir-opt | FileCheck %s --enable-var-scope
 // RUN: rocmlir-gen --arch gfx90a:sramecc+:xnack- --operation attention -current_seq_len=2 -sliding_window_size=1 --causal -return_lse -seq_len_q 1 -seq_len_k 64 -head_dim_qk 32 -head_dim_v 32 -t f32 -pv --apply-bufferization-pipeline=false | rocmlir-opt | FileCheck %s --enable-var-scope --check-prefix=SAFE
+// RUN: rocmlir-gen --arch gfx90a:sramecc+:xnack- --operation attention -g 2 -sliding_window_size=16 -seq_len_q 1 -seq_len_k 64 -head_dim_qk 32 -head_dim_v 32 -t f32 -pv --apply-bufferization-pipeline=false | rocmlir-opt | FileCheck %s --check-prefix=DEFAULT-CURRENT-SEQ-LEN
 
 // CHECK: module attributes {mhal.arch = "[[$ARCH:.*]]"}
 
@@ -57,3 +58,11 @@
 // SAFE: tosa.reciprocal %[[SAFE_SUM]]
 // SAFE: tosa.matmul
 // SAFE: return
+
+// When current_seq_len is omitted, use the last valid key position for every
+// group so tuning-problem keys can be reconstructed by tuningRunner.
+// DEFAULT-CURRENT-SEQ-LEN-LABEL: func.func @rock_attention(
+// DEFAULT-CURRENT-SEQ-LEN-SAME: memref<2xi32>
+// DEFAULT-CURRENT-SEQ-LEN: currentSeqLen = (%{{.*}} : memref<2xi32>)
+// DEFAULT-CURRENT-SEQ-LEN: slidingWindowSize = 16
+// DEFAULT-CURRENT-SEQ-LEN-COUNT-2: arith.constant 63 : i32
diff --git a/mlir/test/rocmlir-gen/options.mlir b/mlir/test/rocmlir-gen/options.mlir
index 4eb69d313ec5..d0b36c5acf42 100644
--- a/mlir/test/rocmlir-gen/options.mlir
+++ b/mlir/test/rocmlir-gen/options.mlir
@@ -37,11 +37,9 @@
 // RUN: not rocmlir-gen --arch %arch --operation attention -t f16 -seq_len_q 256 -seq_len_k 256 -head_dim_qk 32 -head_dim_v 32 -transBias 2>&1 | FileCheck %s --check-prefix=ERR_TRANS_BIAS_WITHOUT_BIAS
 // ERR_TRANS_BIAS_WITHOUT_BIAS: --transBias requires --with-attn-bias
 
-// Sliding-window masking is relative to the KV-cache position.
-// RUN: not rocmlir-gen --arch %arch --operation attention -t f16 -seq_len_q 256 -seq_len_k 256 -head_dim_qk 32 -head_dim_v 32 -sliding_window_size=16 2>&1 | FileCheck %s --check-prefix=ERR_SLIDING_WINDOW
-// ERR_SLIDING_WINDOW: sliding_window_size requires current_seq_len to be set
-
-// A negative value is invalid; zero is the disabled value.
+// A negative sliding_window_size is a user error: the flag's contract is
+// "positive integer, 0 disables", so it must be rejected instead of silently
+// disabling sliding-window masking.
 // RUN: not rocmlir-gen --arch %arch --operation attention -t f16 -seq_len_q 256 -seq_len_k 256 -head_dim_qk 32 -head_dim_v 32 -sliding_window_size=-16 2>&1 | FileCheck %s --check-prefix=ERR_SLIDING_WINDOW_NEG
 // ERR_SLIDING_WINDOW_NEG: sliding_window_size must be non-negative
 
diff --git a/mlir/test/rocmlir-gen/problem-key.mlir b/mlir/test/rocmlir-gen/problem-key.mlir
index d0068d4d742e..0e51faa4a5fa 100644
--- a/mlir/test/rocmlir-gen/problem-key.mlir
+++ b/mlir/test/rocmlir-gen/problem-key.mlir
@@ -38,9 +38,9 @@
 // RUN: rocmlir-gen --arch gfx942 --operation attention -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -t i8 -g 8 | rocmlir-gen --emit-tuning-key - | FileCheck %s --check-prefixes=CHECK_I8_NO_SCALE_BIAS
 // CHECK_I8_NO_SCALE_BIAS: -t i8 {{.*}} -head_dim_v 32 -with-attn-scale false -with-attn-bias false -transBias false
 
-// Sliding-window size affects the generated kernel and is part of its tuning
-// identity. current_seq_len is runtime-only and is intentionally omitted.
-// RUN: rocmlir-gen --arch gfx942 --operation attention -current_seq_len=16 -sliding_window_size 8 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -t f16 -g 1 | rocmlir-gen --emit-tuning-key - | FileCheck %s --check-prefixes=CHECK_SW
+// sliding_window_size is only emitted when set. current_seq_len remains runtime
+// data and defaults to seq_len_k - 1 when this key is reconstructed for tuning.
+// RUN: rocmlir-gen --arch gfx942 --operation attention -sliding_window_size 8 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -t f16 -g 1 | rocmlir-gen --emit-tuning-key - | FileCheck %s --check-prefixes=CHECK_SW
 // CHECK_SW: -t f16 -transQ false -transK false -transV false -transO false -causal false -return_lse false -split_kv 1 -sliding_window_size 8 -num_heads_q 1 -num_heads_kv 1 -g 1 -seq_len_q 256 -seq_len_k 512 -head_dim_qk 64 -head_dim_v 32 -with-attn-scale false -with-attn-bias false -transBias false
 
 // Sliding-window and transposed-bias fields are independent and have stable
diff --git a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp
index 6a711cc44905..824e79613922 100644
--- a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp
+++ b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp
@@ -704,10 +704,13 @@ static llvm::cl::opt<int64_t> splitKV(
 
 static llvm::cl::opt<int64_t> slidingWindowSize(
     "sliding_window_size",
-    llvm::cl::desc("Sliding window attention size. Only the last "
-                   "slidingWindowSize key positions (relative to "
-                   "currentSeqLen) are attended to. Requires current_seq_len."),
-    llvm::cl::value_desc("positive integer"), llvm::cl::init(0));
+    llvm::cl::desc(
+        "Maximum look-back distance from current_seq_len. Includes the current "
+        "KV-cache position, so up to sliding_window_size + 1 key positions are "
+        "attended to. If current_seq_len is omitted, it defaults to "
+        "seq_len_k - 1."),
+    llvm::cl::value_desc("non-negative integer (0 disables)"),
+    llvm::cl::init(0));
 
 static llvm::cl::opt<bool> returnLSE(
     "return_lse",
@@ -1340,6 +1343,9 @@ static LogicalResult detectMissingArguments() {
           << "If split-kv > 1 (flash decoding), we need to return LSE\n";
       return failure();
     }
+    // The flag's contract is "positive integer, 0 disables". A negative value
+    // is a user error that would otherwise slip through silently, since every
+    // downstream use is gated on `slidingWindowSize > 0`.
     if (slidingWindowSize < 0) {
       llvm::errs() << "sliding_window_size must be non-negative\n";
       return failure();
@@ -1350,15 +1356,23 @@ static LogicalResult detectMissingArguments() {
         llvm::errs() << "sliding_window_size must fit in a 32-bit integer\n";
         return failure();
       }
-      if (currentSeqLen.empty()) {
-        llvm::errs()
-            << "sliding_window_size requires current_seq_len to be set\n";
-        return failure();
-      }
+      // The Rock verifier rejects a window larger than the key sequence length
+      // ("slidingWindowSize must not exceed max sequence length"). Reject it
+      // here too so the driver reports a clear error instead of emitting IR
+      // that only fails later in verification.
       if (slidingWindowSize > sequenceLengthK) {
         llvm::errs() << "sliding_window_size must not exceed seq_len_k\n";
         return failure();
       }
+      // Sliding-window masking is defined relative to the KV-cache position, so
+      // the Rock verifier requires currentSeqLen.
+      // currentSeqLen is runtime data and therefore is not part of the tuning
+      // problem key. When a serialized tuning problem is reconstructed by
+      // tuningRunner, use the full-cache position, matching the existing
+      // split-KV default in computeValidSplitKV().
+      if (currentSeqLen.empty())
+        for (int64_t i = 0; i < groupSize; ++i)
+          currentSeqLen.push_back(sequenceLengthK - 1);
     }
   }
 
diff --git a/mlir/utils/performance/perfRunner.py b/mlir/utils/performance/perfRunner.py
index 60b94cb94224..e7d53f88b04f 100644
--- a/mlir/utils/performance/perfRunner.py
+++ b/mlir/utils/performance/perfRunner.py
@@ -1773,7 +1773,8 @@ def __init__(self,
         self.return_lse = return_lse
         self.split_kv = split_kv
         # The window size changes the generated kernel and belongs in its
-        # tuning identity. Runtime sequence positions do not.
+        # tuning identity. Runtime sequence positions do not; rocmlir-gen uses
+        # seq_len_k - 1 for every group when current_seqlen is absent.
         self.sliding_window_size = sliding_window_size
         self.current_seqlen = current_seqlen
 

From ae4c12d43eaf7674bd56770c85d555ece7e57948 Mon Sep 17 00:00:00 2001
From: Umang Yadav <umayadav@amd.com>
Date: Thu, 13 Aug 2026 17:55:56 +0000
Subject: [PATCH 05/10] [NFC] Clarify attention window semantics

The existing wording described currentSeqLen as a sequence length and
slidingWindowSize as a count, which obscured the inclusive endpoint and invited
off-by-one interpretations.

Document currentSeqLen as the last valid zero-based KV-cache position and W as
the maximum look-back distance. The resulting inclusive range contains up to
W + 1 key positions.

(cherry picked from commit c7aed69fc75f39ed98bd4f36cb0534c36d21a1d7)
Co-authored-by: Cursor <cursoragent@cursor.com>
---
 mlir/include/mlir/Dialect/Rock/IR/RockOps.td | 16 ++++++++++++----
 mlir/tools/rocmlir-gen/rocmlir-gen.cpp       | 12 ++++++------
 2 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Rock/IR/RockOps.td b/mlir/include/mlir/Dialect/Rock/IR/RockOps.td
index f0319adf16f8..0189cbbd61cd 100644
--- a/mlir/include/mlir/Dialect/Rock/IR/RockOps.td
+++ b/mlir/include/mlir/Dialect/Rock/IR/RockOps.td
@@ -273,10 +273,15 @@ def Rock_AttentionOp
     - A tensor of shape [G]: per-group/batch offsets, allowing different prefix
       lengths for each sequence in the batch
 
-    If slidingWindowSize is set, we implement sliding window attention where
-    only the last `slidingWindowSize` key positions (relative to currentSeqLen)
-    are attended to. Positions before `max(0, currentSeqLen - slidingWindowSize)`
-    are masked with -inf. This requires currentSeqLen to be set.
+    `currentSeqLen` is the zero-based, inclusive current KV-cache position (the
+    last valid key/value index), not the number of cached entries. For a value
+    `P`, key positions `[0, P]` are valid before other masks are applied.
+
+    If `slidingWindowSize` is set to `W`, it specifies the maximum look-back
+    distance from `currentSeqLen`. Key positions
+    `[max(0, P - W), P]` are attended to, including the current position, so
+    the window contains up to `W + 1` key positions. Earlier positions are
+    masked with -inf. This requires `currentSeqLen` to be set.
 
     LSE (log-sum-exp) is an optional output typically used for flash decoding.
     For flash decoding, you can pass splitKV > 1, the default value is 1, which means flash decoding is disabled. 
@@ -648,6 +653,9 @@ def Rock_GridwiseAttentionAccelOp
   let summary = "Gridwise attention accelerated version";
   let description = [{
     The `rock.gridwise_attention_accel` op computes gridwise attention with acceleration.
+
+    `currentSeqLen` and `slidingWindowSize` use the same inclusive KV-cache
+    position and maximum look-back distance semantics as `rock.attention`.
   }];
   let regions = (region AnyRegion:$preSoftmaxBody);
   let assemblyFormat = [{
diff --git a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp
index 824e79613922..0c8cd958bdcb 100644
--- a/mlir/tools/rocmlir-gen/rocmlir-gen.cpp
+++ b/mlir/tools/rocmlir-gen/rocmlir-gen.cpp
@@ -618,12 +618,12 @@ static llvm::cl::opt<int64_t>
                llvm::cl::desc("number of heads of K,V in attention()"),
                llvm::cl::value_desc("positive integer"), llvm::cl::init(1));
 
-static llvm::cl::list<int64_t>
-    currentSeqLen("current_seq_len",
-                  llvm::cl::desc("List of sequence lengths of K and V (related "
-                                 "to KV-cache) in attention()"),
-                  llvm::cl::value_desc("list of positive integers"),
-                  llvm::cl::CommaSeparated);
+static llvm::cl::list<int64_t> currentSeqLen(
+    "current_seq_len",
+    llvm::cl::desc("List of zero-based, inclusive current KV-cache positions "
+                   "(last valid K/V indices) in attention()"),
+    llvm::cl::value_desc("list of non-negative integers"),
+    llvm::cl::CommaSeparated);
 
 static llvm::cl::opt<int64_t> sequenceLengthQ(
     "seq_len_q", llvm::cl::desc("sequence length of Q in attention()"),

From 500b62cc6e9a5d48731efbe5c252fe562a6d9522 Mon Sep 17 00:00:00 2001
From: Umang Yadav <umayadav@amd.com>
Date: Thu, 13 Aug 2026 17:55:57 +0000
Subject: [PATCH 06/10] Add active sliding-window boundary coverage

Existing attention cases did not force the sliding lower bound to cross an
N-block boundary, so regressions in the N-loop start adjustment could pass while
the window mask remained inactive or inside the first block.

Add random-data F16, BF16, F32, and I8 cases where currentSeqLen - windowSize is
95, inside the [64, 128) block. Use supported attn:v2 configurations so each
case reaches GPU validation instead of failing during lowering.

(cherry picked from commit b4b4f360426a55618ff4ecef6159ec383e2f9145)
Co-authored-by: Cursor <cursoragent@cursor.com>
---
 mlir/test/e2e/PrAttentionBF16.toml | 5 +++++
 mlir/test/e2e/PrAttentionF16.toml  | 5 +++++
 mlir/test/e2e/PrAttentionF32.toml  | 5 +++++
 mlir/test/e2e/PrAttentionI8.toml   | 5 +++++
 4 files changed, 20 insertions(+)

diff --git a/mlir/test/e2e/PrAttentionBF16.toml b/mlir/test/e2e/PrAttentionBF16.toml
index 3176e7d5adb6..61461d28fb49 100644
--- a/mlir/test/e2e/PrAttentionBF16.toml
+++ b/mlir/test/e2e/PrAttentionBF16.toml
@@ -120,6 +120,11 @@ config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -g 3 -num_hea
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -sliding_window_size=64 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
 
+# Active sliding lower mask crossing an N-block boundary. NPerBlock is 64,
+# while currentSeqLen - windowSize is 95, inside the [64, 128) N block.
+[[suite.test]]
+config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=159 -sliding_window_size=64 -g 1 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -perf_config attn:v2:32,32,64,32,32,32,4,1,1,2,1 --with-attn-scale --with-attn-bias"
+
 # GQA + causal + KV Cache batch=3 + return LSE + split-kv (padding)
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 8 -current_seq_len=17,1,32 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
diff --git a/mlir/test/e2e/PrAttentionF16.toml b/mlir/test/e2e/PrAttentionF16.toml
index b1dd7addc99c..66607e95fae7 100644
--- a/mlir/test/e2e/PrAttentionF16.toml
+++ b/mlir/test/e2e/PrAttentionF16.toml
@@ -120,6 +120,11 @@ config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -g 3 -num_hea
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -sliding_window_size=64 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
 
+# Active sliding lower mask crossing an N-block boundary. NPerBlock is 64,
+# while currentSeqLen - windowSize is 95, inside the [64, 128) N block.
+[[suite.test]]
+config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=159 -sliding_window_size=64 -g 1 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -perf_config attn:v2:32,32,64,32,32,32,4,1,1,2,1 --with-attn-scale --with-attn-bias"
+
 # GQA + causal + KV Cache batch=3 + return LSE + split-kv (padding)
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 8 -current_seq_len=17,1,32 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
diff --git a/mlir/test/e2e/PrAttentionF32.toml b/mlir/test/e2e/PrAttentionF32.toml
index 5cd3bda81569..403c6f99148e 100644
--- a/mlir/test/e2e/PrAttentionF32.toml
+++ b/mlir/test/e2e/PrAttentionF32.toml
@@ -92,6 +92,11 @@ config = "-rand 1 -return_lse -split_kv 8 -current_seq_len=17,1,32 -g 3 -num_hea
 #[[suite.test]]
 #config = "-rand 1 -return_lse -split_kv 8 -current_seq_len=17,1,32 -sliding_window_size=64 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
 
+# Active sliding lower mask crossing an N-block boundary. NPerBlock is 64,
+# while currentSeqLen - windowSize is 95, inside the [64, 128) N block.
+[[suite.test]]
+config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=159 -sliding_window_size=64 -g 1 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -perf_config attn:v2:32,32,64,32,32,32,4,1,1,2,1"
+
 # GQA + prefix causal + KV Cache batch=3 + return LSE + split-kv (padding)
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 8 -prefix_offset=18,5,16 -current_seq_len=17,1,32 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
diff --git a/mlir/test/e2e/PrAttentionI8.toml b/mlir/test/e2e/PrAttentionI8.toml
index d3763dbf743f..cb99975c3bc4 100644
--- a/mlir/test/e2e/PrAttentionI8.toml
+++ b/mlir/test/e2e/PrAttentionI8.toml
@@ -97,6 +97,11 @@ config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -g 3 -num_hea
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -sliding_window_size=64 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
 
+# Active sliding lower mask crossing an N-block boundary. NPerBlock is 64,
+# while currentSeqLen - windowSize is 95, inside the [64, 128) N block.
+[[suite.test]]
+config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=159 -sliding_window_size=64 -g 1 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -perf_config attn:v2:32,32,64,16,32,32,8,1,1,2,1 --with-attn-scale --with-attn-bias"
+
 # GQA + causal + KV Cache batch=3 + return LSE + split-kv (padding)
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 8 -current_seq_len=17,1,32 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"

From be805b6249c341576ed83210e852d86eb0578922 Mon Sep 17 00:00:00 2001
From: Umang Yadav <umayadav@amd.com>
Date: Tue, 18 Aug 2026 16:37:29 +0000
Subject: [PATCH 07/10] Reject non-neutral sequence-length transposes

Avoid reconstructing current-sequence-length broadcasts after a transpose changes batch/head ordering. Consolidate shared attention docs and add regression coverage.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 mlir/include/mlir/Dialect/Rock/IR/RockOps.td  |  4 +-
 mlir/lib/Conversion/TosaToRock/TosaToRock.cpp | 65 +++++++++++++------
 .../tosa-to-rock-attention-kvcache.mlir       | 51 +++++++++++++++
 3 files changed, 99 insertions(+), 21 deletions(-)

diff --git a/mlir/include/mlir/Dialect/Rock/IR/RockOps.td b/mlir/include/mlir/Dialect/Rock/IR/RockOps.td
index 0189cbbd61cd..1a92aae6ce3f 100644
--- a/mlir/include/mlir/Dialect/Rock/IR/RockOps.td
+++ b/mlir/include/mlir/Dialect/Rock/IR/RockOps.td
@@ -654,8 +654,8 @@ def Rock_GridwiseAttentionAccelOp
   let description = [{
     The `rock.gridwise_attention_accel` op computes gridwise attention with acceleration.
 
-    `currentSeqLen` and `slidingWindowSize` use the same inclusive KV-cache
-    position and maximum look-back distance semantics as `rock.attention`.
+    See `rock.attention` for additional details on the operands and attributes
+    shared by the two operations.
   }];
   let regions = (region AnyRegion:$preSoftmaxBody);
   let assemblyFormat = [{
diff --git a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
index d76768380dc9..3b932c7bf53a 100644
--- a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
+++ b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
@@ -1965,27 +1965,56 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     return maybeSelect;
   }
 
-  // Helper to verify a value is i32 and traces back to a block argument
-  bool isI32BlockArgument(Value val,
-                          const DenseSet<StringRef> &seqLenSkip) const {
+  // Resolve an i32 sequence-length value to its block argument. A non-trivial
+  // transpose cannot be discarded because it may exchange batch/head values.
+  FailureOr<Value>
+  resolveSeqLenBlockArgument(Value val,
+                             const DenseSet<StringRef> &seqLenSkip) const {
     auto shape = dyn_cast<ShapedType>(val.getType());
     if (!shape || !shape.getElementType().isInteger(32))
-      return false;
+      return failure();
+
+    while (Operation *definingOp = val.getDefiningOp()) {
+      if (!seqLenSkip.contains(definingOp->getName().getStringRef()))
+        break;
+      if (auto transpose = dyn_cast<tosa::TransposeOp>(definingOp)) {
+        ArrayRef<int64_t> inputShape =
+            cast<ShapedType>(transpose.getInput1().getType()).getShape();
+        for (auto [outputDim, inputDim] :
+             llvm::enumerate(transpose.getPerms())) {
+          if (inputDim != static_cast<int32_t>(outputDim) &&
+              (inputShape[outputDim] != 1 || inputShape[inputDim] != 1))
+            return failure();
+        }
+        val = transpose.getInput1();
+      } else if (isa<tosa::MulOp>(definingOp)) {
+        FailureOr<Value> maybeBroadcast = mulBroadcast(val);
+        if (failed(maybeBroadcast))
+          return failure();
+        val = *maybeBroadcast;
+      } else {
+        val = definingOp->getOperand(0);
+      }
+    }
+    if (!isa<BlockArgument>(val))
+      return failure();
+    return val;
+  }
 
-    FailureOr<Value> maybeBlockArg = getValueSkipping(val, seqLenSkip);
-    return succeeded(maybeBlockArg) &&
-           isa<BlockArgument>(maybeBlockArg.value());
+  // Helper to verify a value is i32 and traces back to a block argument.
+  bool isI32BlockArgument(Value val,
+                          const DenseSet<StringRef> &seqLenSkip) const {
+    return succeeded(resolveSeqLenBlockArgument(val, seqLenSkip));
   }
 
   // Returns true when both values resolve to the same currentSeqLen block
   // argument after skipping reshape/broadcast ops.
   bool sameSeqLenBlockArg(Value a, Value b,
                           const DenseSet<StringRef> &seqLenSkip) const {
-    FailureOr<Value> resolvedA = getValueSkipping(a, seqLenSkip);
-    FailureOr<Value> resolvedB = getValueSkipping(b, seqLenSkip);
+    FailureOr<Value> resolvedA = resolveSeqLenBlockArgument(a, seqLenSkip);
+    FailureOr<Value> resolvedB = resolveSeqLenBlockArgument(b, seqLenSkip);
     return succeeded(resolvedA) && succeeded(resolvedB) &&
-           isa<BlockArgument>(resolvedA.value()) &&
-           resolvedA.value() == resolvedB.value();
+           *resolvedA == *resolvedB;
   }
 
   // Helper function to detect select-based causal mask pattern:
@@ -2297,13 +2326,11 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       result.clipMax = maybeClip->clipMax;
     }
 
-    // Resolve through the remaining reshape, transpose, and broadcast chain.
-    // Returning the block argument lets addBroadcastForBlockArg reconstruct
-    // the head broadcast after nested broadcasts have been peeled.
+    // Resolve layout-neutral transforms to the block argument so
+    // addBroadcastForBlockArg can reconstruct the head broadcast.
     FailureOr<Value> maybeCurrentSeqLen =
-        getValueSkipping(seqLenCandidate, seqLenSkip);
-    if (failed(maybeCurrentSeqLen) ||
-        !isI32BlockArgument(*maybeCurrentSeqLen, seqLenSkip))
+        resolveSeqLenBlockArgument(seqLenCandidate, seqLenSkip);
+    if (failed(maybeCurrentSeqLen))
       return failure();
 
     result.seqLen = *maybeCurrentSeqLen;
@@ -2397,8 +2424,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     // An unrelated greater(x - const, col) is not a sliding-window mask. The
     // non-constant operand must resolve to an i32 currentSeqLen block argument.
     FailureOr<Value> maybeSeqLen =
-        getValueSkipping(seqLenCandidate, seqLenSkip);
-    if (failed(maybeSeqLen) || !isI32BlockArgument(*maybeSeqLen, seqLenSkip))
+        resolveSeqLenBlockArgument(seqLenCandidate, seqLenSkip);
+    if (failed(maybeSeqLen))
       return failure();
 
     return SlidingWindowResult{maybeWindowSize.value(), *maybeSeqLen, clipMin,
diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
index a142725c7633..613591e02665 100644
--- a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
+++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
@@ -491,3 +491,54 @@ func.func @mlir_attention_kvcache_sliding_window(%arg0: tensor<1xi32>, %arg1: te
   return %collapsed_9 : tensor<4xf16>
 }
 
+// A non-layout-neutral transpose changes currentSeqLen from per-batch to
+// per-head. Do not discard that mapping and reconstruct a different broadcast.
+// CHECK-LABEL: func @mlir_attention_kvcache_transposed_seqlen
+// CHECK: rock.attention
+// CHECK-NOT: currentSeqLen
+// CHECK: qk = elementwise
+// CHECK: tosa.greater
+// CHECK: tosa.select
+func.func @mlir_attention_kvcache_transposed_seqlen(%arg0: tensor<2xi32>, %arg1: tensor<24xf16>, %arg2: tensor<64xf16>, %arg3: tensor<64xf16>) -> tensor<8xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} {
+  %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %zero = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16>
+  %softmax_ones = "tosa.const"() <{values = dense<1.000000e+00> : tensor<2x2x1x8xf32>}> : () -> tensor<2x2x1x8xf32>
+  %mask_ones = "tosa.const"() <{values = dense<1> : tensor<2x2x1x8xi32>}> : () -> tensor<2x2x1x8xi32>
+  %neg_inf = "tosa.const"() <{values = dense<0xFC00> : tensor<2x2x1x8xf16>}> : () -> tensor<2x2x1x8xf16>
+  %range = "tosa.const"() <{values = dense<[[[[0, 1, 2, 3, 4, 5, 6, 7]]]]> : tensor<1x1x1x8xi32>}> : () -> tensor<1x1x1x8xi32>
+
+  %query = tensor.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [2, 6, 1, 2] : tensor<24xf16> into tensor<2x6x1x2xf16>
+  %query_slice = tensor.extract_slice %query[0, 0, 0, 0] [2, 2, 1, 2] [1, 1, 1, 1] : tensor<2x6x1x2xf16> to tensor<2x2x1x2xf16>
+  %query_collapsed = tensor.collapse_shape %query_slice [[0, 1], [2], [3]] : tensor<2x2x1x2xf16> into tensor<4x1x2xf16>
+  %key = tensor.expand_shape %arg2 [[0, 1, 2, 3]] output_shape [2, 2, 8, 2] : tensor<64xf16> into tensor<2x2x8x2xf16>
+  %key_transposed = tosa.transpose %key {perms = array<i32: 0, 1, 3, 2>} : (tensor<2x2x8x2xf16>) -> tensor<2x2x2x8xf16>
+  %key_collapsed = tensor.collapse_shape %key_transposed [[0, 1], [2], [3]] : tensor<2x2x2x8xf16> into tensor<4x2x8xf16>
+  %qk = tosa.matmul %query_collapsed, %key_collapsed, %zero, %zero : (tensor<4x1x2xf16>, tensor<4x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<4x1x8xf16>
+  %qk_expanded = tensor.expand_shape %qk [[0, 1], [2], [3]] output_shape [2, 2, 1, 8] : tensor<4x1x8xf16> into tensor<2x2x1x8xf16>
+
+  %seq = tensor.expand_shape %arg0 [[0, 1]] output_shape [2, 1] : tensor<2xi32> into tensor<2x1xi32>
+  %seq_transposed = tosa.transpose %seq {perms = array<i32: 1, 0>} : (tensor<2x1xi32>) -> tensor<1x2xi32>
+  %seq_expanded = tensor.expand_shape %seq_transposed [[0], [1, 2, 3]] output_shape [1, 2, 1, 1] : tensor<1x2xi32> into tensor<1x2x1x1xi32>
+  %seq_broadcast = tosa.mul %seq_expanded, %mask_ones, %shift : (tensor<1x2x1x1xi32>, tensor<2x2x1x8xi32>, tensor<1xi8>) -> tensor<2x2x1x8xi32>
+  %range_broadcast = tosa.mul %range, %mask_ones, %shift : (tensor<1x1x1x8xi32>, tensor<2x2x1x8xi32>, tensor<1xi8>) -> tensor<2x2x1x8xi32>
+  %past_end = tosa.greater %range_broadcast, %seq_broadcast : (tensor<2x2x1x8xi32>, tensor<2x2x1x8xi32>) -> tensor<2x2x1x8xi1>
+  %masked = tosa.select %past_end, %neg_inf, %qk_expanded : (tensor<2x2x1x8xi1>, tensor<2x2x1x8xf16>, tensor<2x2x1x8xf16>) -> tensor<2x2x1x8xf16>
+
+  %masked_f32 = tosa.cast %masked : (tensor<2x2x1x8xf16>) -> tensor<2x2x1x8xf32>
+  %max = tosa.reduce_max %masked_f32 {axis = 3 : i32} : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x1xf32>
+  %max_broadcast = tosa.mul %max, %softmax_ones, %shift : (tensor<2x2x1x1xf32>, tensor<2x2x1x8xf32>, tensor<1xi8>) -> tensor<2x2x1x8xf32>
+  %normalized = tosa.sub %masked_f32, %max_broadcast : (tensor<2x2x1x8xf32>, tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf32>
+  %exp = tosa.exp %normalized : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf32>
+  %sum = tosa.reduce_sum %exp {axis = 3 : i32} : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x1xf32>
+  %sum_broadcast = tosa.mul %sum, %softmax_ones, %shift : (tensor<2x2x1x1xf32>, tensor<2x2x1x8xf32>, tensor<1xi8>) -> tensor<2x2x1x8xf32>
+  %reciprocal = tosa.reciprocal %sum_broadcast : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf32>
+  %softmax = tosa.mul %exp, %reciprocal, %shift : (tensor<2x2x1x8xf32>, tensor<2x2x1x8xf32>, tensor<1xi8>) -> tensor<2x2x1x8xf32>
+  %softmax_f16 = tosa.cast %softmax : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf16>
+  %softmax_collapsed = tensor.collapse_shape %softmax_f16 [[0, 1], [2], [3]] : tensor<2x2x1x8xf16> into tensor<4x1x8xf16>
+  %value = tensor.expand_shape %arg3 [[0, 1, 2, 3]] output_shape [2, 2, 8, 2] : tensor<64xf16> into tensor<2x2x8x2xf16>
+  %value_collapsed = tensor.collapse_shape %value [[0, 1], [2], [3]] : tensor<2x2x8x2xf16> into tensor<4x8x2xf16>
+  %output = tosa.matmul %softmax_collapsed, %value_collapsed, %zero, %zero : (tensor<4x1x8xf16>, tensor<4x8x2xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<4x1x2xf16>
+  %result = tensor.collapse_shape %output [[0, 1, 2]] : tensor<4x1x2xf16> into tensor<8xf16>
+  return %result : tensor<8xf16>
+}
+

From 99a70ed692290a61fe30e4a677972d2ee1bb3cee Mon Sep 17 00:00:00 2001
From: Umang Yadav <umayadav@amd.com>
Date: Wed, 19 Aug 2026 17:05:34 +0000
Subject: [PATCH 08/10] Preserve unrepresentable sliding-window masks

Validate window sizes against the key-sequence limit before narrowing and cover split range layouts while correcting block terminology.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 mlir/lib/Conversion/TosaToRock/TosaToRock.cpp |  52 +++--
 ...-to-rock-attention-sliding-window-neg.mlir | 213 ++++++++++++++++++
 mlir/test/e2e/PrAttentionBF16.toml            |   4 +-
 mlir/test/e2e/PrAttentionF16.toml             |   4 +-
 mlir/test/e2e/PrAttentionF32.toml             |   4 +-
 mlir/test/e2e/PrAttentionI8.toml              |   4 +-
 6 files changed, 259 insertions(+), 22 deletions(-)

diff --git a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
index 3b932c7bf53a..77c973555748 100644
--- a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
+++ b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
@@ -54,6 +54,7 @@
 #include "llvm/Support/Debug.h"
 #include "llvm/Support/LogicalResult.h"
 #include "llvm/Support/raw_ostream.h"
+#include <limits>
 #include <numeric>
 #include <tuple>
 #include <utility>
@@ -1750,7 +1751,7 @@ struct AttentionMatcherValues {
   Value currentSeqLen;
   bool isCausal;
   Value prefixOffset;
-  std::optional<int64_t> slidingWindowSize;
+  std::optional<int32_t> slidingWindowSize;
   std::optional<int32_t> seqLenClipMin;
   std::optional<int32_t> seqLenClipMax;
   Type softmaxType;
@@ -2141,7 +2142,7 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     Value inputToContinue; // The value to continue pattern matching with
     Value seqLen;          // The sequence length
     Value prefixOffset;    // The prefix offset value
-    std::optional<int64_t> slidingWindowSize; // The sliding window size
+    std::optional<int32_t> slidingWindowSize; // The sliding window size
     // Clip bounds detected on currentSeqLen during KV-cache pattern matching.
     std::optional<int32_t> seqLenClipMin;
     std::optional<int32_t> seqLenClipMax;
@@ -2339,7 +2340,7 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
   // Result of sliding-window pattern detection.
   struct SlidingWindowResult {
-    int64_t windowSize;
+    int32_t windowSize;
     // The currentSeqLen operand feeding (currentSeqLen - windowSize), resolved
     // through reshape/clip ops so it can be matched against the KV-cache
     // seq-len.
@@ -2354,8 +2355,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
   // greater(add(seqLen, negative_const_offset) * broadcast, col_indices)
   // Returns the window size and validated currentSeqLen operand if successful.
   FailureOr<SlidingWindowResult>
-  trySlidingWindowPattern(Value input,
-                          const DenseSet<StringRef> &seqLenSkip) const {
+  trySlidingWindowPattern(Value input, const DenseSet<StringRef> &seqLenSkip,
+                          int64_t maxSeqLen) const {
     DenseSet<StringRef> expandAndCollapse{
         tensor::CollapseShapeOp::getOperationName(),
         tensor::ExpandShapeOp::getOperationName()};
@@ -2403,6 +2404,13 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     if (failed(maybeWindowSize))
       return failure();
 
+    int64_t windowSize = *maybeWindowSize;
+    // Rock represents the window as an i32 attribute and rejects windows
+    // larger than the key sequence length. Leave those masks explicit.
+    if (windowSize > std::numeric_limits<int32_t>::max() ||
+        windowSize > maxSeqLen)
+      return failure();
+
     // The seq-len operand may be wrapped in a clip (min(max(x, lo), hi)) just
     // like the KV-cache path. Detect it before skipping through the min/max so
     // its bounds can be checked against the KV-cache mask.
@@ -2428,8 +2436,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     if (failed(maybeSeqLen))
       return failure();
 
-    return SlidingWindowResult{maybeWindowSize.value(), *maybeSeqLen, clipMin,
-                               clipMax};
+    return SlidingWindowResult{static_cast<int32_t>(windowSize), *maybeSeqLen,
+                               clipMin, clipMax};
   }
 
   /*
@@ -2593,7 +2601,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
   void analyzeSelectForSeqLenMask(tosa::SelectOp select,
                                   SeqLenMaskResult &result,
                                   const DenseSet<StringRef> &opsToSkip,
-                                  const DenseSet<StringRef> &seqLenSkip) const {
+                                  const DenseSet<StringRef> &seqLenSkip,
+                                  int64_t maxSeqLen) const {
     auto pred = select.getInput1();
     auto maybeGreater = getDefiningOpSkipping<tosa::GreaterOp>(pred, opsToSkip);
     if (failed(maybeGreater))
@@ -2638,7 +2647,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       // would miss the case where the sliding-window mask is seen before the
       // KV-cache mask.
       if (!result.slidingWindowSize) {
-        auto maybeSlidingWindow = trySlidingWindowPattern(input1, seqLenSkip);
+        auto maybeSlidingWindow =
+            trySlidingWindowPattern(input1, seqLenSkip, maxSeqLen);
         if (succeeded(maybeSlidingWindow)) {
           auto slidingWindow = maybeSlidingWindow.value();
           result.slidingWindowSize = slidingWindow.windowSize;
@@ -2651,7 +2661,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     }
   }
 
-  FailureOr<SeqLenMaskResult> getSeqLenMask(Value softmaxInput) const {
+  FailureOr<SeqLenMaskResult> getSeqLenMask(Value softmaxInput,
+                                            int64_t maxSeqLen) const {
     auto maybeSelect = getSelectWithNegInf(softmaxInput);
     if (failed(maybeSelect))
       return failure();
@@ -2677,7 +2688,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
                                    nullptr,         std::nullopt, std::nullopt};
 
     // Analyze the first (outer) select
-    analyzeSelectForSeqLenMask(select, currentResult, opsToSkip, seqLenSkip);
+    analyzeSelectForSeqLenMask(select, currentResult, opsToSkip, seqLenSkip,
+                               maxSeqLen);
 
     // Iteratively peel chained select(mask, -inf, scores) ops to detect
     // separately nested KV-cache, prefix-causal, and sliding-window masks.
@@ -2699,7 +2711,7 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       auto chainedSelect = maybeChainedSelect.value();
       int before = recognizedMaskCount(currentResult);
       analyzeSelectForSeqLenMask(chainedSelect, currentResult, opsToSkip,
-                                 seqLenSkip);
+                                 seqLenSkip, maxSeqLen);
       // Leave an unrecognized or duplicate mask in the elementwise region.
       if (recognizedMaskCount(currentResult) == before)
         break;
@@ -3323,9 +3335,21 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     // or sliding window). Note that non KV-Cache fusions might have
     // tosa.select so, if the checks fail, we just keep going
     Value kvCacheInput, currentSeqLen, prefixOffset;
-    std::optional<int64_t> slidingWindowSize;
+    std::optional<int32_t> slidingWindowSize;
     std::optional<int32_t> seqLenClipMin, seqLenClipMax;
-    auto maybeSeqLenMask = getSeqLenMask(softmaxInput);
+    // Match the Rock verifier's source of truth. The first GEMM's B operand is
+    // the normalized [G, K, N] key tensor, whose trailing dimension is the
+    // maximum key sequence length.
+    ElementwiseRegionFinder<tosa::MatMulOp> softmaxInputFinder;
+    softmaxInputFinder.visit(softmaxInput);
+    FailureOr<tosa::MatMulOp> maybeSourceMatMul =
+        softmaxInputFinder.getFirstGemmBasedOp();
+    if (failed(maybeSourceMatMul))
+      return failure();
+    int64_t maxSeqLen =
+        cast<ShapedType>(maybeSourceMatMul->getB().getType()).getShape().back();
+
+    auto maybeSeqLenMask = getSeqLenMask(softmaxInput, maxSeqLen);
     if (succeeded(maybeSeqLenMask)) {
       auto result = maybeSeqLenMask.value();
       kvCacheInput = result.inputToContinue;
diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
index 85a8953ef1f0..53a422a9d4a8 100644
--- a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
+++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
@@ -289,6 +289,219 @@ func.func @sliding_window_kvcache_negative_clip(%arg0: tensor<1xi32>, %arg1: ten
 
 // -----
 
+// A sliding-window size larger than the key sequence length is rejected by the
+// Rock verifier. Keep the lower mask explicit instead of creating invalid
+// rock.attention IR.
+// CHECK-LABEL: func @sliding_window_exceeds_max_seq_len
+// CHECK: rock.attention
+// CHECK: currentSeqLen =
+// CHECK-NOT: slidingWindowSize
+// CHECK: qk = elementwise
+// CHECK: tosa.add
+// CHECK: tosa.greater
+// CHECK: tosa.select
+func.func @sliding_window_exceeds_max_seq_len(%seq_len: tensor<1xi32>, %queries_flat: tensor<12xf16>, %keys_flat: tensor<32xf16>, %values_flat: tensor<32xf16>) -> tensor<4xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} {
+  %softmax_ones = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x2x1x8xf32>}> : () -> tensor<1x2x1x8xf32>
+  %mask_ones = "tosa.const"() <{values = dense<1> : tensor<1x2x1x8xi8>}> : () -> tensor<1x2x1x8xi8>
+  %window_broadcast_ones = "tosa.const"() <{values = dense<1> : tensor<8x1x1x1xi32>}> : () -> tensor<8x1x1x1xi32>
+  %scale = "tosa.const"() <{values = dense<5.000000e-01> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16>
+  %neg_inf = "tosa.const"() <{values = dense<0xFC00> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16>
+  %zero = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16>
+  %window_offset = "tosa.const"() <{values = dense<-100> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32>
+  %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %kv_broadcast_ones = "tosa.const"() <{values = dense<1> : tensor<1x1x1x8xi32>}> : () -> tensor<1x1x1x8xi32>
+  %columns_1d = "tosa.const"() <{values = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xi32>}> : () -> tensor<8xi32>
+  %columns_4d = arith.constant dense<[[[[0, 1, 2, 3, 4, 5, 6, 7]]]]> : tensor<1x1x1x8xi32>
+  %keys_expanded = tensor.expand_shape %keys_flat [[0, 1, 2, 3]] output_shape [1, 2, 8, 2] : tensor<32xf16> into tensor<1x2x8x2xf16>
+  %queries_expanded = tensor.expand_shape %queries_flat [[0, 1, 2, 3]] output_shape [1, 6, 1, 2] : tensor<12xf16> into tensor<1x6x1x2xf16>
+  %seq_len_4d = tensor.expand_shape %seq_len [[0, 1, 2, 3]] output_shape [1, 1, 1, 1] : tensor<1xi32> into tensor<1x1x1x1xi32>
+  %queries = tensor.extract_slice %queries_expanded[0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1] : tensor<1x6x1x2xf16> to tensor<1x2x1x2xf16>
+  %keys = tosa.transpose %keys_expanded {perms = array<i32: 0, 1, 3, 2>} : (tensor<1x2x8x2xf16>) -> tensor<1x2x2x8xf16>
+  %queries_collapsed = tensor.collapse_shape %queries [[0, 1], [2], [3]] : tensor<1x2x1x2xf16> into tensor<2x1x2xf16>
+  %keys_collapsed = tensor.collapse_shape %keys [[0, 1], [2], [3]] : tensor<1x2x2x8xf16> into tensor<2x2x8xf16>
+  %scores = tosa.matmul %queries_collapsed, %keys_collapsed, %zero, %zero {acc_type = f32} : (tensor<2x1x2xf16>, tensor<2x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x8xf16>
+  %scores_expanded = tensor.expand_shape %scores [[0, 1], [2], [3]] output_shape [1, 2, 1, 8] : tensor<2x1x8xf16> into tensor<1x2x1x8xf16>
+  %scaled_scores = tosa.mul %scores_expanded, %scale, %shift : (tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>, tensor<1xi8>) -> tensor<1x2x1x8xf16>
+  %window_lower = tosa.add %seq_len_4d, %window_offset : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
+  %window_broadcast = tosa.mul %window_lower, %window_broadcast_ones, %shift : (tensor<1x1x1x1xi32>, tensor<8x1x1x1xi32>, tensor<1xi8>) -> tensor<8x1x1x1xi32>
+  %window_flat = tensor.collapse_shape %window_broadcast [[0, 1, 2, 3]] : tensor<8x1x1x1xi32> into tensor<8xi32>
+  %window_pred = tosa.greater %window_flat, %columns_1d : (tensor<8xi32>, tensor<8xi32>) -> tensor<8xi1>
+  %window_pred_i32 = tosa.cast %window_pred : (tensor<8xi1>) -> tensor<8xi32>
+  %window_pred_i8 = tosa.cast %window_pred_i32 : (tensor<8xi32>) -> tensor<8xi8>
+  %window_pred_4d = tensor.expand_shape %window_pred_i8 [[0, 1, 2, 3]] output_shape [1, 1, 1, 8] : tensor<8xi8> into tensor<1x1x1x8xi8>
+  %window_pred_broadcast = tosa.mul %window_pred_4d, %mask_ones, %shift : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
+  %window_mask = tosa.cast %window_pred_broadcast : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
+  %window_masked_scores = tosa.select %window_mask, %neg_inf, %scaled_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %kv_seq_len = tosa.mul %seq_len_4d, %kv_broadcast_ones, %shift : (tensor<1x1x1x1xi32>, tensor<1x1x1x8xi32>, tensor<1xi8>) -> tensor<1x1x1x8xi32>
+  %kv_pred = tosa.greater %columns_4d, %kv_seq_len : (tensor<1x1x1x8xi32>, tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi1>
+  %kv_pred_i32 = tosa.cast %kv_pred : (tensor<1x1x1x8xi1>) -> tensor<1x1x1x8xi32>
+  %kv_pred_i8 = tosa.cast %kv_pred_i32 : (tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi8>
+  %kv_pred_broadcast = tosa.mul %kv_pred_i8, %mask_ones, %shift : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
+  %kv_mask = tosa.cast %kv_pred_broadcast : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
+  %masked_scores = tosa.select %kv_mask, %neg_inf, %window_masked_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %scores_f32 = tosa.cast %masked_scores : (tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf32>
+  %max = tosa.reduce_max %scores_f32 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
+  %max_broadcast = tosa.mul %max, %softmax_ones, %shift : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %normalized = tosa.sub %scores_f32, %max_broadcast : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %exp = tosa.exp %normalized : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %sum = tosa.reduce_sum %exp {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
+  %sum_broadcast = tosa.mul %sum, %softmax_ones, %shift : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %reciprocal = tosa.reciprocal %sum_broadcast : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %softmax = tosa.mul %exp, %reciprocal, %shift : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %softmax_f16 = tosa.cast %softmax : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf16>
+  %softmax_collapsed = tensor.collapse_shape %softmax_f16 [[0, 1], [2], [3]] : tensor<1x2x1x8xf16> into tensor<2x1x8xf16>
+  %values = tensor.expand_shape %values_flat [[0, 1, 2]] output_shape [2, 8, 2] : tensor<32xf16> into tensor<2x8x2xf16>
+  %attention = tosa.matmul %softmax_collapsed, %values, %zero, %zero {acc_type = f32} : (tensor<2x1x8xf16>, tensor<2x8x2xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x2xf16>
+  %attention_expanded = tensor.expand_shape %attention [[0, 1], [2], [3]] output_shape [1, 2, 1, 2] : tensor<2x1x2xf16> into tensor<1x2x1x2xf16>
+  %attention_transposed = tosa.transpose %attention_expanded {perms = array<i32: 0, 2, 1, 3>} : (tensor<1x2x1x2xf16>) -> tensor<1x1x2x2xf16>
+  %result = tensor.collapse_shape %attention_transposed [[0, 1, 2, 3]] : tensor<1x1x2x2xf16> into tensor<4xf16>
+  return %result : tensor<4xf16>
+}
+
+// -----
+
+// The column range may be split across an extra allowed dimension before
+// collapsing to the key-sequence dimension. Validate the window against the
+// first GEMM's key length (8), not the range tensor's trailing dimension (4).
+// CHECK-LABEL: func @sliding_window_uses_key_seq_len
+// CHECK: currentSeqLen =
+// CHECK: slidingWindowSize = 6
+// CHECK: qk = elementwise
+// CHECK-NOT: tosa.greater
+// CHECK-NOT: tosa.select
+// CHECK: rock.yield
+func.func @sliding_window_uses_key_seq_len(%seq_len: tensor<1xi32>, %queries_flat: tensor<12xf16>, %keys_flat: tensor<32xf16>, %values_flat: tensor<32xf16>) -> tensor<4xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} {
+  %softmax_ones = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x2x1x8xf32>}> : () -> tensor<1x2x1x8xf32>
+  %mask_ones = "tosa.const"() <{values = dense<1> : tensor<1x2x1x8xi8>}> : () -> tensor<1x2x1x8xi8>
+  %window_broadcast_ones = "tosa.const"() <{values = dense<1> : tensor<1x1x2x1x4xi32>}> : () -> tensor<1x1x2x1x4xi32>
+  %scale = "tosa.const"() <{values = dense<5.000000e-01> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16>
+  %neg_inf = "tosa.const"() <{values = dense<0xFC00> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16>
+  %zero = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16>
+  %window_offset = "tosa.const"() <{values = dense<-6> : tensor<1x1x1x1x1xi32>}> : () -> tensor<1x1x1x1x1xi32>
+  %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %kv_broadcast_ones = "tosa.const"() <{values = dense<1> : tensor<1x1x1x8xi32>}> : () -> tensor<1x1x1x8xi32>
+  %columns_split = "tosa.const"() <{values = dense<[[[[[0, 1, 2, 3]], [[4, 5, 6, 7]]]]]> : tensor<1x1x2x1x4xi32>}> : () -> tensor<1x1x2x1x4xi32>
+  %columns_4d = arith.constant dense<[[[[0, 1, 2, 3, 4, 5, 6, 7]]]]> : tensor<1x1x1x8xi32>
+  %keys_expanded = tensor.expand_shape %keys_flat [[0, 1, 2, 3]] output_shape [1, 2, 8, 2] : tensor<32xf16> into tensor<1x2x8x2xf16>
+  %queries_expanded = tensor.expand_shape %queries_flat [[0, 1, 2, 3]] output_shape [1, 6, 1, 2] : tensor<12xf16> into tensor<1x6x1x2xf16>
+  %seq_len_4d = tensor.expand_shape %seq_len [[0, 1, 2, 3]] output_shape [1, 1, 1, 1] : tensor<1xi32> into tensor<1x1x1x1xi32>
+  %seq_len_5d = tensor.expand_shape %seq_len [[0, 1, 2, 3, 4]] output_shape [1, 1, 1, 1, 1] : tensor<1xi32> into tensor<1x1x1x1x1xi32>
+  %queries = tensor.extract_slice %queries_expanded[0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1] : tensor<1x6x1x2xf16> to tensor<1x2x1x2xf16>
+  %keys = tosa.transpose %keys_expanded {perms = array<i32: 0, 1, 3, 2>} : (tensor<1x2x8x2xf16>) -> tensor<1x2x2x8xf16>
+  %queries_collapsed = tensor.collapse_shape %queries [[0, 1], [2], [3]] : tensor<1x2x1x2xf16> into tensor<2x1x2xf16>
+  %keys_collapsed = tensor.collapse_shape %keys [[0, 1], [2], [3]] : tensor<1x2x2x8xf16> into tensor<2x2x8xf16>
+  %scores = tosa.matmul %queries_collapsed, %keys_collapsed, %zero, %zero {acc_type = f32} : (tensor<2x1x2xf16>, tensor<2x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x8xf16>
+  %scores_expanded = tensor.expand_shape %scores [[0, 1], [2], [3]] output_shape [1, 2, 1, 8] : tensor<2x1x8xf16> into tensor<1x2x1x8xf16>
+  %scaled_scores = tosa.mul %scores_expanded, %scale, %shift : (tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>, tensor<1xi8>) -> tensor<1x2x1x8xf16>
+  %window_lower = tosa.add %seq_len_5d, %window_offset : (tensor<1x1x1x1x1xi32>, tensor<1x1x1x1x1xi32>) -> tensor<1x1x1x1x1xi32>
+  %window_broadcast = tosa.mul %window_lower, %window_broadcast_ones, %shift : (tensor<1x1x1x1x1xi32>, tensor<1x1x2x1x4xi32>, tensor<1xi8>) -> tensor<1x1x2x1x4xi32>
+  %window_pred_split = tosa.greater %window_broadcast, %columns_split : (tensor<1x1x2x1x4xi32>, tensor<1x1x2x1x4xi32>) -> tensor<1x1x2x1x4xi1>
+  %window_pred = tensor.collapse_shape %window_pred_split [[0, 1, 2, 3, 4]] : tensor<1x1x2x1x4xi1> into tensor<8xi1>
+  %window_pred_i32 = tosa.cast %window_pred : (tensor<8xi1>) -> tensor<8xi32>
+  %window_pred_i8 = tosa.cast %window_pred_i32 : (tensor<8xi32>) -> tensor<8xi8>
+  %window_pred_4d = tensor.expand_shape %window_pred_i8 [[0, 1, 2, 3]] output_shape [1, 1, 1, 8] : tensor<8xi8> into tensor<1x1x1x8xi8>
+  %window_pred_broadcast = tosa.mul %window_pred_4d, %mask_ones, %shift : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
+  %window_mask = tosa.cast %window_pred_broadcast : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
+  %window_masked_scores = tosa.select %window_mask, %neg_inf, %scaled_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %kv_seq_len = tosa.mul %seq_len_4d, %kv_broadcast_ones, %shift : (tensor<1x1x1x1xi32>, tensor<1x1x1x8xi32>, tensor<1xi8>) -> tensor<1x1x1x8xi32>
+  %kv_pred = tosa.greater %columns_4d, %kv_seq_len : (tensor<1x1x1x8xi32>, tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi1>
+  %kv_pred_i32 = tosa.cast %kv_pred : (tensor<1x1x1x8xi1>) -> tensor<1x1x1x8xi32>
+  %kv_pred_i8 = tosa.cast %kv_pred_i32 : (tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi8>
+  %kv_pred_broadcast = tosa.mul %kv_pred_i8, %mask_ones, %shift : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
+  %kv_mask = tosa.cast %kv_pred_broadcast : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
+  %masked_scores = tosa.select %kv_mask, %neg_inf, %window_masked_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %scores_f32 = tosa.cast %masked_scores : (tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf32>
+  %max = tosa.reduce_max %scores_f32 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
+  %max_broadcast = tosa.mul %max, %softmax_ones, %shift : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %normalized = tosa.sub %scores_f32, %max_broadcast : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %exp = tosa.exp %normalized : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %sum = tosa.reduce_sum %exp {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
+  %sum_broadcast = tosa.mul %sum, %softmax_ones, %shift : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %reciprocal = tosa.reciprocal %sum_broadcast : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %softmax = tosa.mul %exp, %reciprocal, %shift : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %softmax_f16 = tosa.cast %softmax : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf16>
+  %softmax_collapsed = tensor.collapse_shape %softmax_f16 [[0, 1], [2], [3]] : tensor<1x2x1x8xf16> into tensor<2x1x8xf16>
+  %values = tensor.expand_shape %values_flat [[0, 1, 2]] output_shape [2, 8, 2] : tensor<32xf16> into tensor<2x8x2xf16>
+  %attention = tosa.matmul %softmax_collapsed, %values, %zero, %zero {acc_type = f32} : (tensor<2x1x8xf16>, tensor<2x8x2xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x2xf16>
+  %attention_expanded = tensor.expand_shape %attention [[0, 1], [2], [3]] output_shape [1, 2, 1, 2] : tensor<2x1x2xf16> into tensor<1x2x1x2xf16>
+  %attention_transposed = tosa.transpose %attention_expanded {perms = array<i32: 0, 2, 1, 3>} : (tensor<1x2x1x2xf16>) -> tensor<1x1x2x2xf16>
+  %result = tensor.collapse_shape %attention_transposed [[0, 1, 2, 3]] : tensor<1x1x2x2xf16> into tensor<4xf16>
+  return %result : tensor<4xf16>
+}
+
+// -----
+
+// Negating INT32_MIN produces a window size that cannot be represented by the
+// i32 slidingWindowSize attribute. Keep the lower mask explicit rather than
+// narrowing it to a negative attribute.
+// CHECK-LABEL: func @sliding_window_int32_min_offset
+// CHECK: rock.attention
+// CHECK: currentSeqLen =
+// CHECK-NOT: slidingWindowSize
+// CHECK: qk = elementwise
+// CHECK: tosa.add
+// CHECK: tosa.greater
+// CHECK: tosa.select
+func.func @sliding_window_int32_min_offset(%seq_len: tensor<1xi32>, %queries_flat: tensor<12xf16>, %keys_flat: tensor<32xf16>, %values_flat: tensor<32xf16>) -> tensor<4xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} {
+  %softmax_ones = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x2x1x8xf32>}> : () -> tensor<1x2x1x8xf32>
+  %mask_ones = "tosa.const"() <{values = dense<1> : tensor<1x2x1x8xi8>}> : () -> tensor<1x2x1x8xi8>
+  %window_broadcast_ones = "tosa.const"() <{values = dense<1> : tensor<8x1x1x1xi32>}> : () -> tensor<8x1x1x1xi32>
+  %scale = "tosa.const"() <{values = dense<5.000000e-01> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16>
+  %neg_inf = "tosa.const"() <{values = dense<0xFC00> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16>
+  %zero = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16>
+  %window_offset = "tosa.const"() <{values = dense<-2147483648> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32>
+  %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %kv_broadcast_ones = "tosa.const"() <{values = dense<1> : tensor<1x1x1x8xi32>}> : () -> tensor<1x1x1x8xi32>
+  %columns_1d = "tosa.const"() <{values = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xi32>}> : () -> tensor<8xi32>
+  %columns_4d = arith.constant dense<[[[[0, 1, 2, 3, 4, 5, 6, 7]]]]> : tensor<1x1x1x8xi32>
+  %keys_expanded = tensor.expand_shape %keys_flat [[0, 1, 2, 3]] output_shape [1, 2, 8, 2] : tensor<32xf16> into tensor<1x2x8x2xf16>
+  %queries_expanded = tensor.expand_shape %queries_flat [[0, 1, 2, 3]] output_shape [1, 6, 1, 2] : tensor<12xf16> into tensor<1x6x1x2xf16>
+  %seq_len_4d = tensor.expand_shape %seq_len [[0, 1, 2, 3]] output_shape [1, 1, 1, 1] : tensor<1xi32> into tensor<1x1x1x1xi32>
+  %queries = tensor.extract_slice %queries_expanded[0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1] : tensor<1x6x1x2xf16> to tensor<1x2x1x2xf16>
+  %keys = tosa.transpose %keys_expanded {perms = array<i32: 0, 1, 3, 2>} : (tensor<1x2x8x2xf16>) -> tensor<1x2x2x8xf16>
+  %queries_collapsed = tensor.collapse_shape %queries [[0, 1], [2], [3]] : tensor<1x2x1x2xf16> into tensor<2x1x2xf16>
+  %keys_collapsed = tensor.collapse_shape %keys [[0, 1], [2], [3]] : tensor<1x2x2x8xf16> into tensor<2x2x8xf16>
+  %scores = tosa.matmul %queries_collapsed, %keys_collapsed, %zero, %zero {acc_type = f32} : (tensor<2x1x2xf16>, tensor<2x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x8xf16>
+  %scores_expanded = tensor.expand_shape %scores [[0, 1], [2], [3]] output_shape [1, 2, 1, 8] : tensor<2x1x8xf16> into tensor<1x2x1x8xf16>
+  %scaled_scores = tosa.mul %scores_expanded, %scale, %shift : (tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>, tensor<1xi8>) -> tensor<1x2x1x8xf16>
+  %window_lower = tosa.add %seq_len_4d, %window_offset : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
+  %window_broadcast = tosa.mul %window_lower, %window_broadcast_ones, %shift : (tensor<1x1x1x1xi32>, tensor<8x1x1x1xi32>, tensor<1xi8>) -> tensor<8x1x1x1xi32>
+  %window_flat = tensor.collapse_shape %window_broadcast [[0, 1, 2, 3]] : tensor<8x1x1x1xi32> into tensor<8xi32>
+  %window_pred = tosa.greater %window_flat, %columns_1d : (tensor<8xi32>, tensor<8xi32>) -> tensor<8xi1>
+  %window_pred_i32 = tosa.cast %window_pred : (tensor<8xi1>) -> tensor<8xi32>
+  %window_pred_i8 = tosa.cast %window_pred_i32 : (tensor<8xi32>) -> tensor<8xi8>
+  %window_pred_4d = tensor.expand_shape %window_pred_i8 [[0, 1, 2, 3]] output_shape [1, 1, 1, 8] : tensor<8xi8> into tensor<1x1x1x8xi8>
+  %window_pred_broadcast = tosa.mul %window_pred_4d, %mask_ones, %shift : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
+  %window_mask = tosa.cast %window_pred_broadcast : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
+  %window_masked_scores = tosa.select %window_mask, %neg_inf, %scaled_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %kv_seq_len = tosa.mul %seq_len_4d, %kv_broadcast_ones, %shift : (tensor<1x1x1x1xi32>, tensor<1x1x1x8xi32>, tensor<1xi8>) -> tensor<1x1x1x8xi32>
+  %kv_pred = tosa.greater %columns_4d, %kv_seq_len : (tensor<1x1x1x8xi32>, tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi1>
+  %kv_pred_i32 = tosa.cast %kv_pred : (tensor<1x1x1x8xi1>) -> tensor<1x1x1x8xi32>
+  %kv_pred_i8 = tosa.cast %kv_pred_i32 : (tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi8>
+  %kv_pred_broadcast = tosa.mul %kv_pred_i8, %mask_ones, %shift : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
+  %kv_mask = tosa.cast %kv_pred_broadcast : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
+  %masked_scores = tosa.select %kv_mask, %neg_inf, %window_masked_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %scores_f32 = tosa.cast %masked_scores : (tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf32>
+  %max = tosa.reduce_max %scores_f32 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
+  %max_broadcast = tosa.mul %max, %softmax_ones, %shift : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %normalized = tosa.sub %scores_f32, %max_broadcast : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %exp = tosa.exp %normalized : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %sum = tosa.reduce_sum %exp {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
+  %sum_broadcast = tosa.mul %sum, %softmax_ones, %shift : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %reciprocal = tosa.reciprocal %sum_broadcast : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %softmax = tosa.mul %exp, %reciprocal, %shift : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %softmax_f16 = tosa.cast %softmax : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf16>
+  %softmax_collapsed = tensor.collapse_shape %softmax_f16 [[0, 1], [2], [3]] : tensor<1x2x1x8xf16> into tensor<2x1x8xf16>
+  %values = tensor.expand_shape %values_flat [[0, 1, 2]] output_shape [2, 8, 2] : tensor<32xf16> into tensor<2x8x2xf16>
+  %attention = tosa.matmul %softmax_collapsed, %values, %zero, %zero {acc_type = f32} : (tensor<2x1x8xf16>, tensor<2x8x2xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x2xf16>
+  %attention_expanded = tensor.expand_shape %attention [[0, 1], [2], [3]] output_shape [1, 2, 1, 2] : tensor<2x1x2xf16> into tensor<1x2x1x2xf16>
+  %attention_transposed = tosa.transpose %attention_expanded {perms = array<i32: 0, 2, 1, 3>} : (tensor<1x2x1x2xf16>) -> tensor<1x1x2x2xf16>
+  %result = tensor.collapse_shape %attention_transposed [[0, 1, 2, 3]] : tensor<1x1x2x2xf16> into tensor<4xf16>
+  return %result : tensor<4xf16>
+}
+
+// -----
+
 // A greater(x - window, col) mask whose x is not currentSeqLen must not be
 // classified as sliding-window attention.
 // CHECK-LABEL: func @not_sliding_window_wrong_operand
diff --git a/mlir/test/e2e/PrAttentionBF16.toml b/mlir/test/e2e/PrAttentionBF16.toml
index 61461d28fb49..afac2eadfb6d 100644
--- a/mlir/test/e2e/PrAttentionBF16.toml
+++ b/mlir/test/e2e/PrAttentionBF16.toml
@@ -120,8 +120,8 @@ config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -g 3 -num_hea
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -sliding_window_size=64 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
 
-# Active sliding lower mask crossing an N-block boundary. NPerBlock is 64,
-# while currentSeqLen - windowSize is 95, inside the [64, 128) N block.
+# Active sliding lower mask in a nonzero key-sequence block. MPerBlockG0 is 32,
+# while currentSeqLen - windowSize is 95, inside the [64, 96) M block.
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=159 -sliding_window_size=64 -g 1 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -perf_config attn:v2:32,32,64,32,32,32,4,1,1,2,1 --with-attn-scale --with-attn-bias"
 
diff --git a/mlir/test/e2e/PrAttentionF16.toml b/mlir/test/e2e/PrAttentionF16.toml
index 66607e95fae7..2c8f7c6ade83 100644
--- a/mlir/test/e2e/PrAttentionF16.toml
+++ b/mlir/test/e2e/PrAttentionF16.toml
@@ -120,8 +120,8 @@ config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -g 3 -num_hea
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -sliding_window_size=64 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
 
-# Active sliding lower mask crossing an N-block boundary. NPerBlock is 64,
-# while currentSeqLen - windowSize is 95, inside the [64, 128) N block.
+# Active sliding lower mask in a nonzero key-sequence block. MPerBlockG0 is 32,
+# while currentSeqLen - windowSize is 95, inside the [64, 96) M block.
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=159 -sliding_window_size=64 -g 1 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -perf_config attn:v2:32,32,64,32,32,32,4,1,1,2,1 --with-attn-scale --with-attn-bias"
 
diff --git a/mlir/test/e2e/PrAttentionF32.toml b/mlir/test/e2e/PrAttentionF32.toml
index 403c6f99148e..647094b3cf29 100644
--- a/mlir/test/e2e/PrAttentionF32.toml
+++ b/mlir/test/e2e/PrAttentionF32.toml
@@ -92,8 +92,8 @@ config = "-rand 1 -return_lse -split_kv 8 -current_seq_len=17,1,32 -g 3 -num_hea
 #[[suite.test]]
 #config = "-rand 1 -return_lse -split_kv 8 -current_seq_len=17,1,32 -sliding_window_size=64 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
 
-# Active sliding lower mask crossing an N-block boundary. NPerBlock is 64,
-# while currentSeqLen - windowSize is 95, inside the [64, 128) N block.
+# Active sliding lower mask in a nonzero key-sequence block. MPerBlockG0 is 32,
+# while currentSeqLen - windowSize is 95, inside the [64, 96) M block.
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=159 -sliding_window_size=64 -g 1 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -perf_config attn:v2:32,32,64,32,32,32,4,1,1,2,1"
 
diff --git a/mlir/test/e2e/PrAttentionI8.toml b/mlir/test/e2e/PrAttentionI8.toml
index cb99975c3bc4..ad5b733c0219 100644
--- a/mlir/test/e2e/PrAttentionI8.toml
+++ b/mlir/test/e2e/PrAttentionI8.toml
@@ -97,8 +97,8 @@ config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -g 3 -num_hea
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=17,1,32 -sliding_window_size=64 -g 3 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -causal --with-attn-scale --with-attn-bias"
 
-# Active sliding lower mask crossing an N-block boundary. NPerBlock is 64,
-# while currentSeqLen - windowSize is 95, inside the [64, 128) N block.
+# Active sliding lower mask in a nonzero key-sequence block. MPerBlockG0 is 32,
+# while currentSeqLen - windowSize is 95, inside the [64, 96) M block.
 [[suite.test]]
 config = "-rand 1 -return_lse -split_kv 4 -current_seq_len=159 -sliding_window_size=64 -g 1 -num_heads_q 4 -num_heads_kv 2 -seq_len_q 1 -seq_len_k 384 -head_dim_qk 64 -head_dim_v 64 -perf_config attn:v2:32,32,64,16,32,32,8,1,1,2,1 --with-attn-scale --with-attn-bias"
 

From 0d41af040261a5ad173070731878e20f7f79a347 Mon Sep 17 00:00:00 2001
From: Umang Yadav <umayadav@amd.com>
Date: Wed, 19 Aug 2026 18:10:30 +0000
Subject: [PATCH 09/10] Preserve one-sided attention clamp semantics

Recognize independent current-sequence clip bounds while keeping unsupported sequence-length and prefix-offset clamps explicit. Verify mask peeling keeps the same QK matmul and document conservative nested-mask fallback.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 mlir/lib/Conversion/TosaToRock/TosaToRock.cpp | 119 ++++++++++--------
 .../tosa-to-rock-attention-kvcache.mlir       |   9 +-
 .../tosa-to-rock-attention-prefix-causal.mlir |  70 ++++++++++-
 ...-to-rock-attention-sliding-window-neg.mlir |  80 +++++++++++-
 4 files changed, 218 insertions(+), 60 deletions(-)

diff --git a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
index 77c973555748..a54912466930 100644
--- a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
+++ b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
@@ -2225,8 +2225,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
   struct ClipResult {
     Value input;
-    int32_t clipMin;
-    int32_t clipMax;
+    std::optional<int32_t> clipMin;
+    std::optional<int32_t> clipMax;
   };
 
   // Peel all multiply-by-one operations used to broadcast a scalar-like value.
@@ -2239,8 +2239,9 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     }
   }
 
-  // Detect min(max(input, clipMin), clipMax), allowing the constant to appear
-  // on either side of each commutative operation.
+  // Detect optional max(input, clipMin) and min(input, clipMax) bounds,
+  // allowing the constant to appear on either side of each commutative
+  // operation.
   FailureOr<ClipResult> tryClipPattern(Value input) const {
     DenseSet<StringRef> expandAndCollapse{
         tensor::CollapseShapeOp::getOperationName(),
@@ -2256,39 +2257,41 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       return attr.getSplatValue<int32_t>();
     };
 
+    Value unclippedInput = input;
+    std::optional<int32_t> clipMin;
+    std::optional<int32_t> clipMax;
+
     auto maybeMin =
         getDefiningOpSkipping<tosa::MinimumOp>(input, expandAndCollapse);
-    if (failed(maybeMin))
-      return failure();
-
-    Value maxCandidate;
-    std::optional<int32_t> clipMax = extractI32Constant(maybeMin->getInput2());
-    if (clipMax) {
-      maxCandidate = maybeMin->getInput1();
-    } else {
-      clipMax = extractI32Constant(maybeMin->getInput1());
-      if (!clipMax)
-        return failure();
-      maxCandidate = maybeMin->getInput2();
+    if (succeeded(maybeMin)) {
+      clipMax = extractI32Constant(maybeMin->getInput2());
+      if (clipMax) {
+        unclippedInput = maybeMin->getInput1();
+      } else {
+        clipMax = extractI32Constant(maybeMin->getInput1());
+        if (!clipMax)
+          return failure();
+        unclippedInput = maybeMin->getInput2();
+      }
     }
 
-    auto maybeMax =
-        getDefiningOpSkipping<tosa::MaximumOp>(maxCandidate, expandAndCollapse);
-    if (failed(maybeMax))
-      return failure();
-
-    Value unclippedInput;
-    std::optional<int32_t> clipMin = extractI32Constant(maybeMax->getInput2());
-    if (clipMin) {
-      unclippedInput = maybeMax->getInput1();
-    } else {
-      clipMin = extractI32Constant(maybeMax->getInput1());
-      if (!clipMin)
-        return failure();
-      unclippedInput = maybeMax->getInput2();
+    auto maybeMax = getDefiningOpSkipping<tosa::MaximumOp>(unclippedInput,
+                                                           expandAndCollapse);
+    if (succeeded(maybeMax)) {
+      clipMin = extractI32Constant(maybeMax->getInput2());
+      if (clipMin) {
+        unclippedInput = maybeMax->getInput1();
+      } else {
+        clipMin = extractI32Constant(maybeMax->getInput1());
+        if (!clipMin)
+          return failure();
+        unclippedInput = maybeMax->getInput2();
+      }
     }
 
-    return ClipResult{unclippedInput, *clipMin, *clipMax};
+    if (!clipMin && !clipMax)
+      return failure();
+    return ClipResult{unclippedInput, clipMin, clipMax};
   }
 
   // Helper to try detecting a KV-cache pattern and an optional clip on its
@@ -2320,7 +2323,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     if (succeeded(maybeClip)) {
       // Rock lowers currentSeqLen masking with unsigned comparisons. Reject a
       // negative clip instead of dropping it while tracing to the block arg.
-      if (maybeClip->clipMin < 0 || maybeClip->clipMax < 0)
+      if ((maybeClip->clipMin && *maybeClip->clipMin < 0) ||
+          (maybeClip->clipMax && *maybeClip->clipMax < 0))
         return failure();
       seqLenCandidate = maybeClip->input;
       result.clipMin = maybeClip->clipMin;
@@ -2411,9 +2415,9 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
         windowSize > maxSeqLen)
       return failure();
 
-    // The seq-len operand may be wrapped in a clip (min(max(x, lo), hi)) just
-    // like the KV-cache path. Detect it before skipping through the min/max so
-    // its bounds can be checked against the KV-cache mask.
+    // The seq-len operand may have lower and/or upper clamp bounds just like
+    // the KV-cache path. Detect them explicitly so they can be checked against
+    // the KV-cache mask; unrecognized clamps must remain in the IR.
     std::optional<int32_t> clipMin;
     std::optional<int32_t> clipMax;
     Value seqLenCandidate = peelBroadcasts(seqLenOperand);
@@ -2422,7 +2426,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     if (succeeded(maybeClip)) {
       // Rock lowers currentSeqLen masking with unsigned comparisons. Reject a
       // negative clip instead of dropping it while tracing to the block arg.
-      if (maybeClip->clipMin < 0 || maybeClip->clipMax < 0)
+      if ((maybeClip->clipMin && *maybeClip->clipMin < 0) ||
+          (maybeClip->clipMax && *maybeClip->clipMax < 0))
         return failure();
       seqLenCandidate = maybeClip->input;
       clipMin = maybeClip->clipMin;
@@ -2674,13 +2679,13 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
                                   tosa::CastOp::getOperationName(),
                                   tosa::MulOp::getOperationName()};
 
-    // Common set used by both pattern detectors
+    // Common set used by both pattern detectors. Min/max clamps are handled
+    // explicitly by tryClipPattern; skipping them here would silently discard
+    // an unrecognized clamp such as one with a non-splat bound.
     DenseSet<StringRef> seqLenSkip{tensor::CollapseShapeOp::getOperationName(),
                                    tensor::ExpandShapeOp::getOperationName(),
                                    tosa::TransposeOp::getOperationName(),
-                                   tosa::MulOp::getOperationName(),
-                                   tosa::MaximumOp::getOperationName(),
-                                   tosa::MinimumOp::getOperationName()};
+                                   tosa::MulOp::getOperationName()};
 
     Value inputToContinue = select.getInput3();
     SeqLenMaskResult currentResult{inputToContinue, nullptr,      nullptr,
@@ -2696,7 +2701,9 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     // Use prefixOffset as the recognition marker for a prefix-causal select
     // (col > row + prefixOffset). A standard causal select (col > row) has no
     // prefixOffset, so it remains in inputToContinue for getCausal() to handle
-    // after the sequence-length masks have been peeled.
+    // after the sequence-length masks have been peeled. Only a contiguous
+    // outer prefix can be peeled: bypassing a recognized mask beneath an
+    // explicit select would require rebuilding the surrounding select chain.
     auto recognizedMaskCount = [](const SeqLenMaskResult &result) {
       return (result.seqLen ? 1 : 0) + (result.prefixOffset ? 1 : 0) +
              (result.slidingWindowSize.has_value() ? 1 : 0);
@@ -2733,9 +2740,9 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       if (!sameSeqLenBlockArg(currentResult.seqLen,
                               currentResult.slidingWindowSeqLen, seqLenSkip))
         return failure();
-      // Both masks must clamp currentSeqLen identically. sameSeqLenBlockArg
-      // only matches the underlying block argument (it skips through the clip
-      // min/max), so a divergent clip would otherwise be silently dropped.
+      // Both masks must clamp currentSeqLen identically. Each recognized clip
+      // has already been resolved to its underlying block argument, so a
+      // divergent pair of bounds would otherwise be silently dropped.
       if (currentResult.seqLenClipMin != currentResult.slidingWindowClipMin ||
           currentResult.seqLenClipMax != currentResult.slidingWindowClipMax)
         return failure();
@@ -3339,13 +3346,19 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     std::optional<int32_t> seqLenClipMin, seqLenClipMax;
     // Match the Rock verifier's source of truth. The first GEMM's B operand is
     // the normalized [G, K, N] key tensor, whose trailing dimension is the
-    // maximum key sequence length.
+    // maximum key sequence length. This traversal must happen before mask
+    // peeling because maxSeqLen determines whether a mask can be peeled, so it
+    // cannot share the post-peeling finder below.
     ElementwiseRegionFinder<tosa::MatMulOp> softmaxInputFinder;
     softmaxInputFinder.visit(softmaxInput);
     FailureOr<tosa::MatMulOp> maybeSourceMatMul =
         softmaxInputFinder.getFirstGemmBasedOp();
-    if (failed(maybeSourceMatMul))
+    if (failed(maybeSourceMatMul)) {
+      LLVM_DEBUG(
+          llvm::dbgs()
+          << "first matmul not found before sequence-length mask analysis\n");
       return failure();
+    }
     int64_t maxSeqLen =
         cast<ShapedType>(maybeSourceMatMul->getB().getType()).getShape().back();
 
@@ -3391,6 +3404,12 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       LLVM_DEBUG(llvm::dbgs() << "first matmul not found\n");
       return failure();
     }
+    if (*maybeFirstMatMul != *maybeSourceMatMul) {
+      LLVM_DEBUG(
+          llvm::dbgs()
+          << "first matmul changed after sequence-length mask analysis\n");
+      return failure();
+    }
 
     TypedValue<TensorType> matC = maybeFirstMatMul.value().getOutput();
     ArrayRef<int64_t> shapeC = matC.getType().getShape();
@@ -3529,10 +3548,10 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     prepareBlockArgTensor(currentSeqLen);
     prepareBlockArgTensor(prefixOffset);
 
-    // Apply seqLen clip if detected during KV-cache pattern matching.
-    // The original model may have clip(arg, lo, hi) on currentSeqLen which
-    // was traced through to reach the block argument. The clip is a property
-    // of currentSeqLen itself, used by all masks (KV-cache, sliding window).
+    // Apply any seqLen clip bounds detected during KV-cache pattern matching.
+    // The original model's clamp was traced through to reach the block
+    // argument. It is a property of currentSeqLen itself, used by all masks
+    // (KV-cache, sliding window).
     if (currentSeqLen && (attentionMatcherValues.seqLenClipMin.has_value() ||
                           attentionMatcherValues.seqLenClipMax.has_value())) {
       auto seqLenType = cast<RankedTensorType>(currentSeqLen.getType());
diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
index 613591e02665..816e67522f5f 100644
--- a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
+++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
@@ -415,8 +415,8 @@ func.func @mlir_causal_attention_nokvcache_wrongrange(%arg0: tensor<24576xf16>,
 
 // CHECK-LABEL:func @mlir_attention_kvcache_sliding_window
 // CHECK: %[[MAX:.*]] = tosa.maximum
-// CHECK: %[[CLIP:.*]] = tosa.minimum %[[MAX]], {{.*}} : (tensor<2xi32>, tensor<2xi32>) -> tensor<2xi32>
-// CHECK: currentSeqLen = (%[[CLIP]] : tensor<2xi32>)
+// CHECK-NOT: tosa.minimum
+// CHECK: currentSeqLen = (%[[MAX]] : tensor<2xi32>)
 // CHECK: slidingWindowSize = 3
 func.func @mlir_attention_kvcache_sliding_window(%arg0: tensor<1xi32>, %arg1: tensor<12xf16>, %arg2: tensor<32xf16>, %arg3: tensor<32xf16>) -> tensor<4xf16> attributes {rock.kernel = "mixr"} {
   %0 = "tosa.const"() <{values = dense<4> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32>
@@ -447,7 +447,6 @@ func.func @mlir_attention_kvcache_sliding_window(%arg0: tensor<1xi32>, %arg1: te
   %cst = arith.constant dense<[[[[0, 1, 2, 3, 4, 5, 6, 7]]]]> : tensor<1x1x1x8xi32>
   %expanded_1 = tensor.expand_shape %arg0 [[0, 1, 2, 3]] output_shape [1, 1, 1, 1] : tensor<1xi32> into tensor<1x1x1x1xi32>
   %23 = tosa.maximum %expanded_1, %0 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
-  %24 = tosa.minimum %23, %0 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
   %extracted_slice = tensor.extract_slice %expanded_0[0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1] : tensor<1x6x1x2xf16> to tensor<1x2x1x2xf16>
   %25 = tosa.transpose %expanded {perms = array<i32: 0, 1, 3, 2>} : (tensor<1x2x8x2xf16>) -> tensor<1x2x2x8xf16>
   %collapsed = tensor.collapse_shape %extracted_slice [[0, 1], [2], [3]] : tensor<1x2x1x2xf16> into tensor<2x1x2xf16>
@@ -455,7 +454,7 @@ func.func @mlir_attention_kvcache_sliding_window(%arg0: tensor<1xi32>, %arg1: te
   %26 = tosa.matmul %collapsed, %collapsed_2, %11, %11 {acc_type = f32} : (tensor<2x1x2xf16>, tensor<2x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x8xf16>
   %expanded_3 = tensor.expand_shape %26 [[0, 1], [2], [3]] output_shape [1, 2, 1, 8] : tensor<2x1x8xf16> into tensor<1x2x1x8xf16>
   %27 = tosa.mul %expanded_3, %8, %17 : (tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>, tensor<1xi8>) -> tensor<1x2x1x8xf16>
-  %28 = tosa.add %24, %16 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
+  %28 = tosa.add %23, %16 : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
   %29 = tosa.mul %28, %7, %17 : (tensor<1x1x1x1xi32>, tensor<8x1x1x1xi32>, tensor<1xi8>) -> tensor<8x1x1x1xi32>
   %collapsed_4 = tensor.collapse_shape %29 [[0, 1, 2, 3]] : tensor<8x1x1x1xi32> into tensor<8xi32>
   %30 = tosa.greater %collapsed_4, %20 : (tensor<8xi32>, tensor<8xi32>) -> tensor<8xi1>
@@ -465,7 +464,7 @@ func.func @mlir_attention_kvcache_sliding_window(%arg0: tensor<1xi32>, %arg1: te
   %33 = tosa.mul %expanded_5, %5, %17 : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
   %34 = tosa.cast %33 : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
   %35 = tosa.select %34, %9, %27 : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
-  %36 = tosa.mul %24, %18, %17 : (tensor<1x1x1x1xi32>, tensor<1x1x1x8xi32>, tensor<1xi8>) -> tensor<1x1x1x8xi32>
+  %36 = tosa.mul %23, %18, %17 : (tensor<1x1x1x1xi32>, tensor<1x1x1x8xi32>, tensor<1xi8>) -> tensor<1x1x1x8xi32>
   %37 = tosa.greater %cst, %36 : (tensor<1x1x1x8xi32>, tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi1>
   %38 = tosa.cast %37 : (tensor<1x1x1x8xi1>) -> tensor<1x1x1x8xi32>
   %39 = tosa.cast %38 : (tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi8>
diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-prefix-causal.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-prefix-causal.mlir
index d225cd066fbb..0f17e464a8a9 100644
--- a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-prefix-causal.mlir
+++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-prefix-causal.mlir
@@ -1,4 +1,4 @@
-// RUN: rocmlir-opt --tosa-to-rock %s | FileCheck %s
+// RUN: rocmlir-opt --tosa-to-rock %s -verify-diagnostics | FileCheck %s
 
 module {
   // CHECK-LABEL: func @mlir_attention
@@ -69,5 +69,73 @@ module {
     %collapsed_8 = tensor.collapse_shape %48 [[0, 1, 2, 3]] : tensor<1x4x14x64xf16> into tensor<3584xf16>
     return %collapsed_8 : tensor<3584xf16>
   }
+
+  // A clamped prefix offset is not recognized by the prefix-causal matcher.
+  // Keep the mask explicit instead of dropping the clamp while resolving the
+  // offset to its block argument.
+  // CHECK-LABEL: func @mlir_attention_clamped_prefix_offset
+  // CHECK: rock.attention
+  // CHECK-NOT: prefixOffset
+  // CHECK-NOT: causal
+  // CHECK: qk = elementwise
+  // CHECK: tosa.minimum
+  // CHECK: tosa.add
+  // CHECK: tosa.greater
+  // CHECK: tosa.select
+  func.func @mlir_attention_clamped_prefix_offset(%arg0: tensor<1xi32>, %arg1: tensor<4608xf16>, %arg2: tensor<2048xf16>, %arg3: tensor<14336xf16>) -> tensor<3584xf16> attributes {rock.kernel} {
+    %columns = "tosa.const"() <{values = dense<[[0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]]> : tensor<1x16xi32>}> : () -> tensor<1x16xi32>
+    %softmax_ones = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x14x4x16xf32>}> : () -> tensor<1x14x4x16xf32>
+    %neg_inf = "tosa.const"() <{values = dense<0xFC00> : tensor<1x14x4x16xf16>}> : () -> tensor<1x14x4x16xf16>
+    %zero = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16>
+    %query_ones = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x2x7x64x16xf16>}> : () -> tensor<1x2x7x64x16xf16>
+    %mask_ones = "tosa.const"() <{values = dense<1> : tensor<1x14x4x16xi8>}> : () -> tensor<1x14x4x16xi8>
+    %row_broadcast_ones = "tosa.const"() <{values = dense<1> : tensor<4x1xi32>}> : () -> tensor<4x1xi32>
+    %scale = "tosa.const"() <{values = dense<1.250000e-01> : tensor<1x14x4x16xf16>}> : () -> tensor<1x14x4x16xf16>
+    %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+    %column_broadcast_ones = "tosa.const"() <{values = dense<1> : tensor<4x16xi32>}> : () -> tensor<4x16xi32>
+    %rows = "tosa.const"() <{values = dense<[[0], [1], [2], [3]]> : tensor<4x1xi32>}> : () -> tensor<4x1xi32>
+    %clip_max = "tosa.const"() <{values = dense<8> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
+    %keys_expanded = tensor.expand_shape %arg2 [[0, 1, 2, 3, 4]] output_shape [1, 2, 1, 16, 64] : tensor<2048xf16> into tensor<1x2x1x16x64xf16>
+    %queries_expanded = tensor.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [1, 4, 18, 64] : tensor<4608xf16> into tensor<1x4x18x64xf16>
+    %queries_transposed = tosa.transpose %queries_expanded {perms = array<i32: 0, 2, 1, 3>} : (tensor<1x4x18x64xf16>) -> tensor<1x18x4x64xf16>
+    %offset = tensor.expand_shape %arg0 [[0, 1]] output_shape [1, 1] : tensor<1xi32> into tensor<1x1xi32>
+    %clipped_offset = tosa.minimum %offset, %clip_max : (tensor<1x1xi32>, tensor<1x1xi32>) -> tensor<1x1xi32>
+    %columns_4d = tosa.mul %columns, %column_broadcast_ones, %shift : (tensor<1x16xi32>, tensor<4x16xi32>, tensor<1xi8>) -> tensor<4x16xi32>
+    %offset_4d = tosa.mul %clipped_offset, %row_broadcast_ones, %shift : (tensor<1x1xi32>, tensor<4x1xi32>, tensor<1xi8>) -> tensor<4x1xi32>
+    %row_bound = tosa.add %offset_4d, %rows : (tensor<4x1xi32>, tensor<4x1xi32>) -> tensor<4x1xi32>
+    %row_bound_broadcast = tosa.mul %row_bound, %column_broadcast_ones, %shift : (tensor<4x1xi32>, tensor<4x16xi32>, tensor<1xi8>) -> tensor<4x16xi32>
+    %mask_pred = tosa.greater %columns_4d, %row_bound_broadcast : (tensor<4x16xi32>, tensor<4x16xi32>) -> tensor<4x16xi1>
+    %mask_i32 = tosa.cast %mask_pred : (tensor<4x16xi1>) -> tensor<4x16xi32>
+    %mask_i8 = tosa.cast %mask_i32 : (tensor<4x16xi32>) -> tensor<4x16xi8>
+    %mask_expanded = tensor.expand_shape %mask_i8 [[0, 1, 2], [3]] output_shape [1, 1, 4, 16] : tensor<4x16xi8> into tensor<1x1x4x16xi8>
+    %mask_broadcast = tosa.mul %mask_expanded, %mask_ones, %shift : (tensor<1x1x4x16xi8>, tensor<1x14x4x16xi8>, tensor<1xi8>) -> tensor<1x14x4x16xi8>
+    %queries = tensor.extract_slice %queries_transposed[0, 0, 0, 0] [1, 14, 4, 64] [1, 1, 1, 1] : tensor<1x18x4x64xf16> to tensor<1x14x4x64xf16>
+    %keys_transposed = tosa.transpose %keys_expanded {perms = array<i32: 0, 1, 2, 4, 3>} : (tensor<1x2x1x16x64xf16>) -> tensor<1x2x1x64x16xf16>
+    %keys_broadcast = tosa.mul %keys_transposed, %query_ones, %shift : (tensor<1x2x1x64x16xf16>, tensor<1x2x7x64x16xf16>, tensor<1xi8>) -> tensor<1x2x7x64x16xf16>
+    %queries_collapsed = tensor.collapse_shape %queries [[0, 1], [2], [3]] : tensor<1x14x4x64xf16> into tensor<14x4x64xf16>
+    %keys_collapsed = tensor.collapse_shape %keys_broadcast [[0, 1, 2], [3], [4]] : tensor<1x2x7x64x16xf16> into tensor<14x64x16xf16>
+    %scores = tosa.matmul %queries_collapsed, %keys_collapsed, %zero, %zero {acc_type = f32} : (tensor<14x4x64xf16>, tensor<14x64x16xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<14x4x16xf16>
+    %scores_expanded = tensor.expand_shape %scores [[0, 1], [2], [3]] output_shape [1, 14, 4, 16] : tensor<14x4x16xf16> into tensor<1x14x4x16xf16>
+    %scaled_scores = tosa.mul %scores_expanded, %scale, %shift : (tensor<1x14x4x16xf16>, tensor<1x14x4x16xf16>, tensor<1xi8>) -> tensor<1x14x4x16xf16>
+    %mask = tosa.cast %mask_broadcast : (tensor<1x14x4x16xi8>) -> tensor<1x14x4x16xi1>
+    %masked_scores = tosa.select %mask, %neg_inf, %scaled_scores : (tensor<1x14x4x16xi1>, tensor<1x14x4x16xf16>, tensor<1x14x4x16xf16>) -> tensor<1x14x4x16xf16>
+    %scores_f32 = tosa.cast %masked_scores : (tensor<1x14x4x16xf16>) -> tensor<1x14x4x16xf32>
+    %max = tosa.reduce_max %scores_f32 {axis = 3 : i32} : (tensor<1x14x4x16xf32>) -> tensor<1x14x4x1xf32>
+    %max_broadcast = tosa.mul %max, %softmax_ones, %shift : (tensor<1x14x4x1xf32>, tensor<1x14x4x16xf32>, tensor<1xi8>) -> tensor<1x14x4x16xf32>
+    %normalized = tosa.sub %scores_f32, %max_broadcast : (tensor<1x14x4x16xf32>, tensor<1x14x4x16xf32>) -> tensor<1x14x4x16xf32>
+    %exp = tosa.exp %normalized : (tensor<1x14x4x16xf32>) -> tensor<1x14x4x16xf32>
+    %sum = tosa.reduce_sum %exp {axis = 3 : i32} : (tensor<1x14x4x16xf32>) -> tensor<1x14x4x1xf32>
+    %sum_broadcast = tosa.mul %sum, %softmax_ones, %shift : (tensor<1x14x4x1xf32>, tensor<1x14x4x16xf32>, tensor<1xi8>) -> tensor<1x14x4x16xf32>
+    %reciprocal = tosa.reciprocal %sum_broadcast : (tensor<1x14x4x16xf32>) -> tensor<1x14x4x16xf32>
+    %softmax = tosa.mul %exp, %reciprocal, %shift : (tensor<1x14x4x16xf32>, tensor<1x14x4x16xf32>, tensor<1xi8>) -> tensor<1x14x4x16xf32>
+    %softmax_f16 = tosa.cast %softmax : (tensor<1x14x4x16xf32>) -> tensor<1x14x4x16xf16>
+    %softmax_collapsed = tensor.collapse_shape %softmax_f16 [[0, 1], [2], [3]] : tensor<1x14x4x16xf16> into tensor<14x4x16xf16>
+    %values = tensor.expand_shape %arg3 [[0, 1, 2]] output_shape [14, 16, 64] : tensor<14336xf16> into tensor<14x16x64xf16>
+    %attention = tosa.matmul %softmax_collapsed, %values, %zero, %zero {acc_type = f32} : (tensor<14x4x16xf16>, tensor<14x16x64xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<14x4x64xf16>
+    %attention_expanded = tensor.expand_shape %attention [[0, 1], [2], [3]] output_shape [1, 14, 4, 64] : tensor<14x4x64xf16> into tensor<1x14x4x64xf16>
+    %attention_transposed = tosa.transpose %attention_expanded {perms = array<i32: 0, 2, 1, 3>} : (tensor<1x14x4x64xf16>) -> tensor<1x4x14x64xf16>
+    %result = tensor.collapse_shape %attention_transposed [[0, 1, 2, 3]] : tensor<1x4x14x64xf16> into tensor<3584xf16>
+    return %result : tensor<3584xf16>
+  }
 }
 
diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
index 53a422a9d4a8..f1c00baad409 100644
--- a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
+++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-sliding-window-neg.mlir
@@ -289,6 +289,76 @@ func.func @sliding_window_kvcache_negative_clip(%arg0: tensor<1xi32>, %arg1: ten
 
 // -----
 
+// A one-sided upper clamp is representable by rematerializing the minimum on
+// currentSeqLen. Fold both masks without introducing a lower clamp.
+// CHECK-LABEL: func @sliding_window_kvcache_one_sided_clip
+// CHECK-NOT: tosa.maximum
+// CHECK: %[[CLIP:.*]] = tosa.minimum
+// CHECK: rock.attention
+// CHECK: currentSeqLen = (%[[CLIP]] : tensor<2xi32>)
+// CHECK: slidingWindowSize = 3
+func.func @sliding_window_kvcache_one_sided_clip(%arg0: tensor<1xi32>, %arg1: tensor<12xf16>, %arg2: tensor<32xf16>, %arg3: tensor<32xf16>) -> tensor<4xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} {
+  %clip_max = "tosa.const"() <{values = dense<4> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32>
+  %softmax_ones = "tosa.const"() <{values = dense<1.000000e+00> : tensor<1x2x1x8xf32>}> : () -> tensor<1x2x1x8xf32>
+  %mask_ones = "tosa.const"() <{values = dense<1> : tensor<1x2x1x8xi8>}> : () -> tensor<1x2x1x8xi8>
+  %window_broadcast_ones = "tosa.const"() <{values = dense<1> : tensor<8x1x1x1xi32>}> : () -> tensor<8x1x1x1xi32>
+  %scale = "tosa.const"() <{values = dense<5.000000e-01> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16>
+  %neg_inf = "tosa.const"() <{values = dense<0xFC00> : tensor<1x2x1x8xf16>}> : () -> tensor<1x2x1x8xf16>
+  %zero = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16>
+  %window_offset = "tosa.const"() <{values = dense<-3> : tensor<1x1x1x1xi32>}> : () -> tensor<1x1x1x1xi32>
+  %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %kv_broadcast_ones = "tosa.const"() <{values = dense<1> : tensor<1x1x1x8xi32>}> : () -> tensor<1x1x1x8xi32>
+  %columns_1d = "tosa.const"() <{values = dense<[0, 1, 2, 3, 4, 5, 6, 7]> : tensor<8xi32>}> : () -> tensor<8xi32>
+  %columns_4d = arith.constant dense<[[[[0, 1, 2, 3, 4, 5, 6, 7]]]]> : tensor<1x1x1x8xi32>
+  %keys_expanded = tensor.expand_shape %arg2 [[0, 1, 2, 3]] output_shape [1, 2, 8, 2] : tensor<32xf16> into tensor<1x2x8x2xf16>
+  %queries_expanded = tensor.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [1, 6, 1, 2] : tensor<12xf16> into tensor<1x6x1x2xf16>
+  %seq_len = tensor.expand_shape %arg0 [[0, 1, 2, 3]] output_shape [1, 1, 1, 1] : tensor<1xi32> into tensor<1x1x1x1xi32>
+  %clipped_seq_len = tosa.minimum %seq_len, %clip_max : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
+  %queries = tensor.extract_slice %queries_expanded[0, 0, 0, 0] [1, 2, 1, 2] [1, 1, 1, 1] : tensor<1x6x1x2xf16> to tensor<1x2x1x2xf16>
+  %keys = tosa.transpose %keys_expanded {perms = array<i32: 0, 1, 3, 2>} : (tensor<1x2x8x2xf16>) -> tensor<1x2x2x8xf16>
+  %queries_collapsed = tensor.collapse_shape %queries [[0, 1], [2], [3]] : tensor<1x2x1x2xf16> into tensor<2x1x2xf16>
+  %keys_collapsed = tensor.collapse_shape %keys [[0, 1], [2], [3]] : tensor<1x2x2x8xf16> into tensor<2x2x8xf16>
+  %scores = tosa.matmul %queries_collapsed, %keys_collapsed, %zero, %zero {acc_type = f32} : (tensor<2x1x2xf16>, tensor<2x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x8xf16>
+  %scores_expanded = tensor.expand_shape %scores [[0, 1], [2], [3]] output_shape [1, 2, 1, 8] : tensor<2x1x8xf16> into tensor<1x2x1x8xf16>
+  %scaled_scores = tosa.mul %scores_expanded, %scale, %shift : (tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>, tensor<1xi8>) -> tensor<1x2x1x8xf16>
+  %kv_seq_len = tosa.mul %clipped_seq_len, %kv_broadcast_ones, %shift : (tensor<1x1x1x1xi32>, tensor<1x1x1x8xi32>, tensor<1xi8>) -> tensor<1x1x1x8xi32>
+  %kv_pred = tosa.greater %columns_4d, %kv_seq_len : (tensor<1x1x1x8xi32>, tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi1>
+  %kv_pred_i32 = tosa.cast %kv_pred : (tensor<1x1x1x8xi1>) -> tensor<1x1x1x8xi32>
+  %kv_pred_i8 = tosa.cast %kv_pred_i32 : (tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi8>
+  %kv_pred_broadcast = tosa.mul %kv_pred_i8, %mask_ones, %shift : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
+  %kv_mask = tosa.cast %kv_pred_broadcast : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
+  %kv_masked_scores = tosa.select %kv_mask, %neg_inf, %scaled_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %window_lower = tosa.add %clipped_seq_len, %window_offset : (tensor<1x1x1x1xi32>, tensor<1x1x1x1xi32>) -> tensor<1x1x1x1xi32>
+  %window_broadcast = tosa.mul %window_lower, %window_broadcast_ones, %shift : (tensor<1x1x1x1xi32>, tensor<8x1x1x1xi32>, tensor<1xi8>) -> tensor<8x1x1x1xi32>
+  %window_flat = tensor.collapse_shape %window_broadcast [[0, 1, 2, 3]] : tensor<8x1x1x1xi32> into tensor<8xi32>
+  %window_pred = tosa.greater %window_flat, %columns_1d : (tensor<8xi32>, tensor<8xi32>) -> tensor<8xi1>
+  %window_pred_i32 = tosa.cast %window_pred : (tensor<8xi1>) -> tensor<8xi32>
+  %window_pred_i8 = tosa.cast %window_pred_i32 : (tensor<8xi32>) -> tensor<8xi8>
+  %window_pred_4d = tensor.expand_shape %window_pred_i8 [[0, 1, 2, 3]] output_shape [1, 1, 1, 8] : tensor<8xi8> into tensor<1x1x1x8xi8>
+  %window_pred_broadcast = tosa.mul %window_pred_4d, %mask_ones, %shift : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
+  %window_mask = tosa.cast %window_pred_broadcast : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
+  %masked_scores = tosa.select %window_mask, %neg_inf, %kv_masked_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %scores_f32 = tosa.cast %masked_scores : (tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf32>
+  %max = tosa.reduce_max %scores_f32 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
+  %max_broadcast = tosa.mul %max, %softmax_ones, %shift : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %normalized = tosa.sub %scores_f32, %max_broadcast : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %exp = tosa.exp %normalized : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %sum = tosa.reduce_sum %exp {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
+  %sum_broadcast = tosa.mul %sum, %softmax_ones, %shift : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %reciprocal = tosa.reciprocal %sum_broadcast : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf32>
+  %softmax = tosa.mul %exp, %reciprocal, %shift : (tensor<1x2x1x8xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>
+  %softmax_f16 = tosa.cast %softmax : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x8xf16>
+  %softmax_collapsed = tensor.collapse_shape %softmax_f16 [[0, 1], [2], [3]] : tensor<1x2x1x8xf16> into tensor<2x1x8xf16>
+  %values = tensor.expand_shape %arg3 [[0, 1, 2]] output_shape [2, 8, 2] : tensor<32xf16> into tensor<2x8x2xf16>
+  %attention = tosa.matmul %softmax_collapsed, %values, %zero, %zero {acc_type = f32} : (tensor<2x1x8xf16>, tensor<2x8x2xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<2x1x2xf16>
+  %attention_expanded = tensor.expand_shape %attention [[0, 1], [2], [3]] output_shape [1, 2, 1, 2] : tensor<2x1x2xf16> into tensor<1x2x1x2xf16>
+  %attention_transposed = tosa.transpose %attention_expanded {perms = array<i32: 0, 2, 1, 3>} : (tensor<1x2x1x2xf16>) -> tensor<1x1x2x2xf16>
+  %result = tensor.collapse_shape %attention_transposed [[0, 1, 2, 3]] : tensor<1x1x2x2xf16> into tensor<4xf16>
+  return %result : tensor<4xf16>
+}
+
+// -----
+
 // A sliding-window size larger than the key sequence length is rejected by the
 // Rock verifier. Keep the lower mask explicit instead of creating invalid
 // rock.attention IR.
@@ -433,10 +503,12 @@ func.func @sliding_window_uses_key_seq_len(%seq_len: tensor<1xi32>, %queries_fla
 
 // Negating INT32_MIN produces a window size that cannot be represented by the
 // i32 slidingWindowSize attribute. Keep the lower mask explicit rather than
-// narrowing it to a negative attribute.
+// narrowing it to a negative attribute. Since that unsupported mask is the
+// outer select, keep the nested KV-cache mask explicit too; peeling it alone
+// would require rebuilding the surrounding select chain.
 // CHECK-LABEL: func @sliding_window_int32_min_offset
 // CHECK: rock.attention
-// CHECK: currentSeqLen =
+// CHECK-NOT: currentSeqLen
 // CHECK-NOT: slidingWindowSize
 // CHECK: qk = elementwise
 // CHECK: tosa.add
@@ -473,14 +545,14 @@ func.func @sliding_window_int32_min_offset(%seq_len: tensor<1xi32>, %queries_fla
   %window_pred_4d = tensor.expand_shape %window_pred_i8 [[0, 1, 2, 3]] output_shape [1, 1, 1, 8] : tensor<8xi8> into tensor<1x1x1x8xi8>
   %window_pred_broadcast = tosa.mul %window_pred_4d, %mask_ones, %shift : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
   %window_mask = tosa.cast %window_pred_broadcast : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
-  %window_masked_scores = tosa.select %window_mask, %neg_inf, %scaled_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
   %kv_seq_len = tosa.mul %seq_len_4d, %kv_broadcast_ones, %shift : (tensor<1x1x1x1xi32>, tensor<1x1x1x8xi32>, tensor<1xi8>) -> tensor<1x1x1x8xi32>
   %kv_pred = tosa.greater %columns_4d, %kv_seq_len : (tensor<1x1x1x8xi32>, tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi1>
   %kv_pred_i32 = tosa.cast %kv_pred : (tensor<1x1x1x8xi1>) -> tensor<1x1x1x8xi32>
   %kv_pred_i8 = tosa.cast %kv_pred_i32 : (tensor<1x1x1x8xi32>) -> tensor<1x1x1x8xi8>
   %kv_pred_broadcast = tosa.mul %kv_pred_i8, %mask_ones, %shift : (tensor<1x1x1x8xi8>, tensor<1x2x1x8xi8>, tensor<1xi8>) -> tensor<1x2x1x8xi8>
   %kv_mask = tosa.cast %kv_pred_broadcast : (tensor<1x2x1x8xi8>) -> tensor<1x2x1x8xi1>
-  %masked_scores = tosa.select %kv_mask, %neg_inf, %window_masked_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %kv_masked_scores = tosa.select %kv_mask, %neg_inf, %scaled_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
+  %masked_scores = tosa.select %window_mask, %neg_inf, %kv_masked_scores : (tensor<1x2x1x8xi1>, tensor<1x2x1x8xf16>, tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf16>
   %scores_f32 = tosa.cast %masked_scores : (tensor<1x2x1x8xf16>) -> tensor<1x2x1x8xf32>
   %max = tosa.reduce_max %scores_f32 {axis = 3 : i32} : (tensor<1x2x1x8xf32>) -> tensor<1x2x1x1xf32>
   %max_broadcast = tosa.mul %max, %softmax_ones, %shift : (tensor<1x2x1x1xf32>, tensor<1x2x1x8xf32>, tensor<1xi8>) -> tensor<1x2x1x8xf32>

From 0cf8d9dfd451f2e63d867f4825298f5f0a1e2d60 Mon Sep 17 00:00:00 2001
From: Umang Yadav <umayadav@amd.com>
Date: Wed, 19 Aug 2026 21:15:36 +0000
Subject: [PATCH 10/10] Reject unsafe sequence-length folding

Preserve batch/head provenance and leave out-of-range clips explicit to prevent silent value reordering and oversized key traversal.

Co-authored-by: Cursor <cursoragent@cursor.com>
---
 mlir/lib/Conversion/TosaToRock/TosaToRock.cpp | 124 +++++++++++-------
 .../tosa-to-rock-attention-kvcache.mlir       | 103 +++++++++++++++
 ...-sliding-window-kvcache-prefix-causal.mlir |   3 +-
 3 files changed, 181 insertions(+), 49 deletions(-)

diff --git a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
index a54912466930..a434a1a8d099 100644
--- a/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
+++ b/mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
@@ -1967,10 +1967,11 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
   }
 
   // Resolve an i32 sequence-length value to its block argument. A non-trivial
-  // transpose cannot be discarded because it may exchange batch/head values.
+  // transpose or reshape is discarded only when it preserves the group-value
+  // interpretation expected by attention lowering.
   FailureOr<Value>
-  resolveSeqLenBlockArgument(Value val,
-                             const DenseSet<StringRef> &seqLenSkip) const {
+  resolveSeqLenBlockArgument(Value val, const DenseSet<StringRef> &seqLenSkip,
+                             int64_t expectedNumGroups) const {
     auto shape = dyn_cast<ShapedType>(val.getType());
     if (!shape || !shape.getElementType().isInteger(32))
       return failure();
@@ -1993,8 +1994,25 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
         if (failed(maybeBroadcast))
           return failure();
         val = *maybeBroadcast;
+      } else if (auto expand = dyn_cast<tensor::ExpandShapeOp>(definingOp)) {
+        ArrayRef<int64_t> inputShape = expand.getSrcType().getShape();
+        ArrayRef<int64_t> outputShape = expand.getResultType().getShape();
+        // A non-scalar rank-1 input shorter than the group count is interpreted
+        // as per-batch when lowering reconstructs its broadcast. Only discard
+        // its expansion when the value remains on the leading axis; otherwise
+        // per-head values would be silently reconstructed as per-batch values.
+        if (inputShape.size() == 1 && inputShape.front() != 1 &&
+            inputShape.front() != expectedNumGroups &&
+            (outputShape.empty() || outputShape.front() != inputShape.front() ||
+             !llvm::all_of(outputShape.drop_front(),
+                           [](int64_t dim) { return dim == 1; })))
+          return failure();
+        val = expand.getSrc();
+      } else if (auto collapse =
+                     dyn_cast<tensor::CollapseShapeOp>(definingOp)) {
+        val = collapse.getSrc();
       } else {
-        val = definingOp->getOperand(0);
+        return failure();
       }
     }
     if (!isa<BlockArgument>(val))
@@ -2002,18 +2020,15 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     return val;
   }
 
-  // Helper to verify a value is i32 and traces back to a block argument.
-  bool isI32BlockArgument(Value val,
-                          const DenseSet<StringRef> &seqLenSkip) const {
-    return succeeded(resolveSeqLenBlockArgument(val, seqLenSkip));
-  }
-
   // Returns true when both values resolve to the same currentSeqLen block
   // argument after skipping reshape/broadcast ops.
   bool sameSeqLenBlockArg(Value a, Value b,
-                          const DenseSet<StringRef> &seqLenSkip) const {
-    FailureOr<Value> resolvedA = resolveSeqLenBlockArgument(a, seqLenSkip);
-    FailureOr<Value> resolvedB = resolveSeqLenBlockArgument(b, seqLenSkip);
+                          const DenseSet<StringRef> &seqLenSkip,
+                          int64_t expectedNumGroups) const {
+    FailureOr<Value> resolvedA =
+        resolveSeqLenBlockArgument(a, seqLenSkip, expectedNumGroups);
+    FailureOr<Value> resolvedB =
+        resolveSeqLenBlockArgument(b, seqLenSkip, expectedNumGroups);
     return succeeded(resolvedA) && succeeded(resolvedB) &&
            *resolvedA == *resolvedB;
   }
@@ -2164,9 +2179,9 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
   // Helper to try detecting prefix causal pattern: add(row_indices, offset)
   // Returns the offset value if successful
-  FailureOr<Value>
-  tryPrefixCausalPattern(Value input,
-                         const DenseSet<StringRef> &seqLenSkip) const {
+  FailureOr<Value> tryPrefixCausalPattern(Value input,
+                                          const DenseSet<StringRef> &seqLenSkip,
+                                          int64_t expectedNumGroups) const {
     DenseSet<StringRef> expandAndCollapse{
         tensor::CollapseShapeOp::getOperationName(),
         tensor::ExpandShapeOp::getOperationName()};
@@ -2203,16 +2218,11 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       maybeOffset = offset;
 
     FailureOr<Value> maybeOffsetUnwrapped =
-        getValueSkipping(maybeOffset.value(), expandAndCollapse);
+        resolveSeqLenBlockArgument(*maybeOffset, seqLenSkip, expectedNumGroups);
     if (failed(maybeOffsetUnwrapped))
       return failure();
 
     Value unwrappedOffset = maybeOffsetUnwrapped.value();
-
-    // Verify offset is i32 and traces back to a block argument
-    if (!isI32BlockArgument(unwrappedOffset, seqLenSkip))
-      return failure();
-
     return unwrappedOffset;
   }
 
@@ -2294,10 +2304,21 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
     return ClipResult{unclippedInput, clipMin, clipMax};
   }
 
+  bool hasValidSeqLenClipBounds(const ClipResult &clip,
+                                int64_t maxSeqLen) const {
+    if (maxSeqLen <= 0)
+      return false;
+    auto isOutOfRange = [maxSeqLen](std::optional<int32_t> bound) {
+      return bound && (*bound < 0 || static_cast<int64_t>(*bound) >= maxSeqLen);
+    };
+    return !isOutOfRange(clip.clipMin) && !isOutOfRange(clip.clipMax);
+  }
+
   // Helper to try detecting a KV-cache pattern and an optional clip on its
   // sequence length.
   FailureOr<KVCacheResult>
-  tryKVCachePattern(Value input, const DenseSet<StringRef> &seqLenSkip) const {
+  tryKVCachePattern(Value input, const DenseSet<StringRef> &seqLenSkip,
+                    int64_t maxSeqLen, int64_t expectedNumGroups) const {
     FailureOr<Value> maybeNonOne = mulBroadcast(input);
     if (failed(maybeNonOne))
       return failure();
@@ -2321,10 +2342,10 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
     auto maybeClip = tryClipPattern(seqLenCandidate);
     if (succeeded(maybeClip)) {
-      // Rock lowers currentSeqLen masking with unsigned comparisons. Reject a
-      // negative clip instead of dropping it while tracing to the block arg.
-      if ((maybeClip->clipMin && *maybeClip->clipMin < 0) ||
-          (maybeClip->clipMax && *maybeClip->clipMax < 0))
+      // Rock interprets currentSeqLen as an unsigned, inclusive key position
+      // and does not cap the noncausal traversal. Leave a clip explicit unless
+      // every bound is a valid key position.
+      if (!hasValidSeqLenClipBounds(*maybeClip, maxSeqLen))
         return failure();
       seqLenCandidate = maybeClip->input;
       result.clipMin = maybeClip->clipMin;
@@ -2333,8 +2354,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
     // Resolve layout-neutral transforms to the block argument so
     // addBroadcastForBlockArg can reconstruct the head broadcast.
-    FailureOr<Value> maybeCurrentSeqLen =
-        resolveSeqLenBlockArgument(seqLenCandidate, seqLenSkip);
+    FailureOr<Value> maybeCurrentSeqLen = resolveSeqLenBlockArgument(
+        seqLenCandidate, seqLenSkip, expectedNumGroups);
     if (failed(maybeCurrentSeqLen))
       return failure();
 
@@ -2360,7 +2381,7 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
   // Returns the window size and validated currentSeqLen operand if successful.
   FailureOr<SlidingWindowResult>
   trySlidingWindowPattern(Value input, const DenseSet<StringRef> &seqLenSkip,
-                          int64_t maxSeqLen) const {
+                          int64_t maxSeqLen, int64_t expectedNumGroups) const {
     DenseSet<StringRef> expandAndCollapse{
         tensor::CollapseShapeOp::getOperationName(),
         tensor::ExpandShapeOp::getOperationName()};
@@ -2424,10 +2445,9 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
     auto maybeClip = tryClipPattern(seqLenCandidate);
     if (succeeded(maybeClip)) {
-      // Rock lowers currentSeqLen masking with unsigned comparisons. Reject a
-      // negative clip instead of dropping it while tracing to the block arg.
-      if ((maybeClip->clipMin && *maybeClip->clipMin < 0) ||
-          (maybeClip->clipMax && *maybeClip->clipMax < 0))
+      // Keep clip validation consistent with the KV-cache mask. Invalid bounds
+      // must remain in the explicit select chain.
+      if (!hasValidSeqLenClipBounds(*maybeClip, maxSeqLen))
         return failure();
       seqLenCandidate = maybeClip->input;
       clipMin = maybeClip->clipMin;
@@ -2436,8 +2456,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
     // An unrelated greater(x - const, col) is not a sliding-window mask. The
     // non-constant operand must resolve to an i32 currentSeqLen block argument.
-    FailureOr<Value> maybeSeqLen =
-        resolveSeqLenBlockArgument(seqLenCandidate, seqLenSkip);
+    FailureOr<Value> maybeSeqLen = resolveSeqLenBlockArgument(
+        seqLenCandidate, seqLenSkip, expectedNumGroups);
     if (failed(maybeSeqLen))
       return failure();
 
@@ -2607,7 +2627,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
                                   SeqLenMaskResult &result,
                                   const DenseSet<StringRef> &opsToSkip,
                                   const DenseSet<StringRef> &seqLenSkip,
-                                  int64_t maxSeqLen) const {
+                                  int64_t maxSeqLen,
+                                  int64_t expectedNumGroups) const {
     auto pred = select.getInput1();
     auto maybeGreater = getDefiningOpSkipping<tosa::GreaterOp>(pred, opsToSkip);
     if (failed(maybeGreater))
@@ -2622,7 +2643,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
       // Try KV-cache pattern (scalar seqLen) if not already found
       if (!result.seqLen) {
-        auto maybeKVCache = tryKVCachePattern(input2, seqLenSkip);
+        auto maybeKVCache =
+            tryKVCachePattern(input2, seqLenSkip, maxSeqLen, expectedNumGroups);
         if (succeeded(maybeKVCache)) {
           auto kvResult = maybeKVCache.value();
           result.seqLen = kvResult.seqLen;
@@ -2633,7 +2655,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
       // Try prefix causal pattern (row_indices + offset) if not already found
       if (!result.prefixOffset) {
-        auto maybePrefixCausal = tryPrefixCausalPattern(input2, seqLenSkip);
+        auto maybePrefixCausal =
+            tryPrefixCausalPattern(input2, seqLenSkip, expectedNumGroups);
         if (succeeded(maybePrefixCausal)) {
           result.prefixOffset = maybePrefixCausal.value();
         }
@@ -2652,8 +2675,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       // would miss the case where the sliding-window mask is seen before the
       // KV-cache mask.
       if (!result.slidingWindowSize) {
-        auto maybeSlidingWindow =
-            trySlidingWindowPattern(input1, seqLenSkip, maxSeqLen);
+        auto maybeSlidingWindow = trySlidingWindowPattern(
+            input1, seqLenSkip, maxSeqLen, expectedNumGroups);
         if (succeeded(maybeSlidingWindow)) {
           auto slidingWindow = maybeSlidingWindow.value();
           result.slidingWindowSize = slidingWindow.windowSize;
@@ -2667,7 +2690,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
   }
 
   FailureOr<SeqLenMaskResult> getSeqLenMask(Value softmaxInput,
-                                            int64_t maxSeqLen) const {
+                                            int64_t maxSeqLen,
+                                            int64_t expectedNumGroups) const {
     auto maybeSelect = getSelectWithNegInf(softmaxInput);
     if (failed(maybeSelect))
       return failure();
@@ -2694,7 +2718,7 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
 
     // Analyze the first (outer) select
     analyzeSelectForSeqLenMask(select, currentResult, opsToSkip, seqLenSkip,
-                               maxSeqLen);
+                               maxSeqLen, expectedNumGroups);
 
     // Iteratively peel chained select(mask, -inf, scores) ops to detect
     // separately nested KV-cache, prefix-causal, and sliding-window masks.
@@ -2718,7 +2742,7 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       auto chainedSelect = maybeChainedSelect.value();
       int before = recognizedMaskCount(currentResult);
       analyzeSelectForSeqLenMask(chainedSelect, currentResult, opsToSkip,
-                                 seqLenSkip, maxSeqLen);
+                                 seqLenSkip, maxSeqLen, expectedNumGroups);
       // Leave an unrecognized or duplicate mask in the elementwise region.
       if (recognizedMaskCount(currentResult) == before)
         break;
@@ -2738,7 +2762,8 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
       // as the KV-cache mask; otherwise the two masks disagree on the sequence
       // length and cannot be folded into an op with one currentSeqLen.
       if (!sameSeqLenBlockArg(currentResult.seqLen,
-                              currentResult.slidingWindowSeqLen, seqLenSkip))
+                              currentResult.slidingWindowSeqLen, seqLenSkip,
+                              expectedNumGroups))
         return failure();
       // Both masks must clamp currentSeqLen identically. Each recognized clip
       // has already been resolved to its underlying block argument, so a
@@ -3359,10 +3384,13 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
           << "first matmul not found before sequence-length mask analysis\n");
       return failure();
     }
-    int64_t maxSeqLen =
-        cast<ShapedType>(maybeSourceMatMul->getB().getType()).getShape().back();
+    ArrayRef<int64_t> keyShape =
+        cast<ShapedType>(maybeSourceMatMul->getB().getType()).getShape();
+    int64_t expectedNumGroups = keyShape.front();
+    int64_t maxSeqLen = keyShape.back();
 
-    auto maybeSeqLenMask = getSeqLenMask(softmaxInput, maxSeqLen);
+    auto maybeSeqLenMask =
+        getSeqLenMask(softmaxInput, maxSeqLen, expectedNumGroups);
     if (succeeded(maybeSeqLenMask)) {
       auto result = maybeSeqLenMask.value();
       kvCacheInput = result.inputToContinue;
diff --git a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
index 816e67522f5f..838a8428d369 100644
--- a/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
+++ b/mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-kvcache.mlir
@@ -541,3 +541,106 @@ func.func @mlir_attention_kvcache_transposed_seqlen(%arg0: tensor<2xi32>, %arg1:
   return %result : tensor<8xf16>
 }
 
+// Expanding a rank-1 sequence tensor onto the head axis makes it per-head, not
+// per-batch. Do not strip that expansion and reconstruct a different broadcast.
+// CHECK-LABEL: func @mlir_attention_kvcache_per_head_broadcast_seqlen
+// CHECK: rock.attention
+// CHECK-NOT: currentSeqLen
+// CHECK: qk = elementwise
+// CHECK: tosa.greater
+// CHECK: tosa.select
+func.func @mlir_attention_kvcache_per_head_broadcast_seqlen(%arg0: tensor<2xi32>, %arg1: tensor<24xf16>, %arg2: tensor<64xf16>, %arg3: tensor<64xf16>) -> tensor<8xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} {
+  %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %zero = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16>
+  %softmax_ones = "tosa.const"() <{values = dense<1.000000e+00> : tensor<2x2x1x8xf32>}> : () -> tensor<2x2x1x8xf32>
+  %mask_ones = "tosa.const"() <{values = dense<1> : tensor<2x2x1x8xi32>}> : () -> tensor<2x2x1x8xi32>
+  %head_ones = "tosa.const"() <{values = dense<1> : tensor<2x2x1x1xi32>}> : () -> tensor<2x2x1x1xi32>
+  %neg_inf = "tosa.const"() <{values = dense<0xFC00> : tensor<2x2x1x8xf16>}> : () -> tensor<2x2x1x8xf16>
+  %range = "tosa.const"() <{values = dense<[[[[0, 1, 2, 3, 4, 5, 6, 7]]]]> : tensor<1x1x1x8xi32>}> : () -> tensor<1x1x1x8xi32>
+
+  %query = tensor.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [2, 6, 1, 2] : tensor<24xf16> into tensor<2x6x1x2xf16>
+  %query_slice = tensor.extract_slice %query[0, 0, 0, 0] [2, 2, 1, 2] [1, 1, 1, 1] : tensor<2x6x1x2xf16> to tensor<2x2x1x2xf16>
+  %query_collapsed = tensor.collapse_shape %query_slice [[0, 1], [2], [3]] : tensor<2x2x1x2xf16> into tensor<4x1x2xf16>
+  %key = tensor.expand_shape %arg2 [[0, 1, 2, 3]] output_shape [2, 2, 8, 2] : tensor<64xf16> into tensor<2x2x8x2xf16>
+  %key_transposed = tosa.transpose %key {perms = array<i32: 0, 1, 3, 2>} : (tensor<2x2x8x2xf16>) -> tensor<2x2x2x8xf16>
+  %key_collapsed = tensor.collapse_shape %key_transposed [[0, 1], [2], [3]] : tensor<2x2x2x8xf16> into tensor<4x2x8xf16>
+  %qk = tosa.matmul %query_collapsed, %key_collapsed, %zero, %zero : (tensor<4x1x2xf16>, tensor<4x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<4x1x8xf16>
+  %qk_expanded = tensor.expand_shape %qk [[0, 1], [2], [3]] output_shape [2, 2, 1, 8] : tensor<4x1x8xf16> into tensor<2x2x1x8xf16>
+
+  %seq_head = tensor.expand_shape %arg0 [[0, 1, 2, 3]] output_shape [1, 2, 1, 1] : tensor<2xi32> into tensor<1x2x1x1xi32>
+  %seq_groups = tosa.mul %seq_head, %head_ones, %shift : (tensor<1x2x1x1xi32>, tensor<2x2x1x1xi32>, tensor<1xi8>) -> tensor<2x2x1x1xi32>
+  %seq_broadcast = tosa.mul %seq_groups, %mask_ones, %shift : (tensor<2x2x1x1xi32>, tensor<2x2x1x8xi32>, tensor<1xi8>) -> tensor<2x2x1x8xi32>
+  %range_broadcast = tosa.mul %range, %mask_ones, %shift : (tensor<1x1x1x8xi32>, tensor<2x2x1x8xi32>, tensor<1xi8>) -> tensor<2x2x1x8xi32>
+  %past_end = tosa.greater %range_broadcast, %seq_broadcast : (tensor<2x2x1x8xi32>, tensor<2x2x1x8xi32>) -> tensor<2x2x1x8xi1>
+  %masked = tosa.select %past_end, %neg_inf, %qk_expanded : (tensor<2x2x1x8xi1>, tensor<2x2x1x8xf16>, tensor<2x2x1x8xf16>) -> tensor<2x2x1x8xf16>
+
+  %masked_f32 = tosa.cast %masked : (tensor<2x2x1x8xf16>) -> tensor<2x2x1x8xf32>
+  %max = tosa.reduce_max %masked_f32 {axis = 3 : i32} : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x1xf32>
+  %max_broadcast = tosa.mul %max, %softmax_ones, %shift : (tensor<2x2x1x1xf32>, tensor<2x2x1x8xf32>, tensor<1xi8>) -> tensor<2x2x1x8xf32>
+  %normalized = tosa.sub %masked_f32, %max_broadcast : (tensor<2x2x1x8xf32>, tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf32>
+  %exp = tosa.exp %normalized : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf32>
+  %sum = tosa.reduce_sum %exp {axis = 3 : i32} : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x1xf32>
+  %sum_broadcast = tosa.mul %sum, %softmax_ones, %shift : (tensor<2x2x1x1xf32>, tensor<2x2x1x8xf32>, tensor<1xi8>) -> tensor<2x2x1x8xf32>
+  %reciprocal = tosa.reciprocal %sum_broadcast : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf32>
+  %softmax = tosa.mul %exp, %reciprocal, %shift : (tensor<2x2x1x8xf32>, tensor<2x2x1x8xf32>, tensor<1xi8>) -> tensor<2x2x1x8xf32>
+  %softmax_f16 = tosa.cast %softmax : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf16>
+  %softmax_collapsed = tensor.collapse_shape %softmax_f16 [[0, 1], [2], [3]] : tensor<2x2x1x8xf16> into tensor<4x1x8xf16>
+  %value = tensor.expand_shape %arg3 [[0, 1, 2, 3]] output_shape [2, 2, 8, 2] : tensor<64xf16> into tensor<2x2x8x2xf16>
+  %value_collapsed = tensor.collapse_shape %value [[0, 1], [2], [3]] : tensor<2x2x8x2xf16> into tensor<4x8x2xf16>
+  %output = tosa.matmul %softmax_collapsed, %value_collapsed, %zero, %zero : (tensor<4x1x8xf16>, tensor<4x8x2xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<4x1x2xf16>
+  %result = tensor.collapse_shape %output [[0, 1, 2]] : tensor<4x1x2xf16> into tensor<8xf16>
+  return %result : tensor<8xf16>
+}
+
+// currentSeqLen is an inclusive key position, so a clamp bound equal to the
+// key length is outside its valid range. Keep this mask explicit.
+// CHECK-LABEL: func @mlir_attention_kvcache_out_of_range_lower_clip
+// CHECK: rock.attention
+// CHECK-NOT: currentSeqLen
+// CHECK: qk = elementwise
+// CHECK: tosa.maximum
+// CHECK: tosa.greater
+// CHECK: tosa.select
+func.func @mlir_attention_kvcache_out_of_range_lower_clip(%arg0: tensor<2xi32>, %arg1: tensor<24xf16>, %arg2: tensor<64xf16>, %arg3: tensor<64xf16>) -> tensor<8xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} {
+  %shift = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
+  %zero = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16>
+  %softmax_ones = "tosa.const"() <{values = dense<1.000000e+00> : tensor<2x2x1x8xf32>}> : () -> tensor<2x2x1x8xf32>
+  %mask_ones = "tosa.const"() <{values = dense<1> : tensor<2x2x1x8xi32>}> : () -> tensor<2x2x1x8xi32>
+  %clip_min = "tosa.const"() <{values = dense<8> : tensor<2x1x1x1xi32>}> : () -> tensor<2x1x1x1xi32>
+  %neg_inf = "tosa.const"() <{values = dense<0xFC00> : tensor<2x2x1x8xf16>}> : () -> tensor<2x2x1x8xf16>
+  %range = "tosa.const"() <{values = dense<[[[[0, 1, 2, 3, 4, 5, 6, 7]]]]> : tensor<1x1x1x8xi32>}> : () -> tensor<1x1x1x8xi32>
+
+  %query = tensor.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [2, 6, 1, 2] : tensor<24xf16> into tensor<2x6x1x2xf16>
+  %query_slice = tensor.extract_slice %query[0, 0, 0, 0] [2, 2, 1, 2] [1, 1, 1, 1] : tensor<2x6x1x2xf16> to tensor<2x2x1x2xf16>
+  %query_collapsed = tensor.collapse_shape %query_slice [[0, 1], [2], [3]] : tensor<2x2x1x2xf16> into tensor<4x1x2xf16>
+  %key = tensor.expand_shape %arg2 [[0, 1, 2, 3]] output_shape [2, 2, 8, 2] : tensor<64xf16> into tensor<2x2x8x2xf16>
+  %key_transposed = tosa.transpose %key {perms = array<i32: 0, 1, 3, 2>} : (tensor<2x2x8x2xf16>) -> tensor<2x2x2x8xf16>
+  %key_collapsed = tensor.collapse_shape %key_transposed [[0, 1], [2], [3]] : tensor<2x2x2x8xf16> into tensor<4x2x8xf16>
+  %qk = tosa.matmul %query_collapsed, %key_collapsed, %zero, %zero : (tensor<4x1x2xf16>, tensor<4x2x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<4x1x8xf16>
+  %qk_expanded = tensor.expand_shape %qk [[0, 1], [2], [3]] output_shape [2, 2, 1, 8] : tensor<4x1x8xf16> into tensor<2x2x1x8xf16>
+
+  %seq = tensor.expand_shape %arg0 [[0, 1, 2, 3]] output_shape [2, 1, 1, 1] : tensor<2xi32> into tensor<2x1x1x1xi32>
+  %clipped_seq = tosa.maximum %seq, %clip_min : (tensor<2x1x1x1xi32>, tensor<2x1x1x1xi32>) -> tensor<2x1x1x1xi32>
+  %seq_broadcast = tosa.mul %clipped_seq, %mask_ones, %shift : (tensor<2x1x1x1xi32>, tensor<2x2x1x8xi32>, tensor<1xi8>) -> tensor<2x2x1x8xi32>
+  %range_broadcast = tosa.mul %range, %mask_ones, %shift : (tensor<1x1x1x8xi32>, tensor<2x2x1x8xi32>, tensor<1xi8>) -> tensor<2x2x1x8xi32>
+  %past_end = tosa.greater %range_broadcast, %seq_broadcast : (tensor<2x2x1x8xi32>, tensor<2x2x1x8xi32>) -> tensor<2x2x1x8xi1>
+  %masked = tosa.select %past_end, %neg_inf, %qk_expanded : (tensor<2x2x1x8xi1>, tensor<2x2x1x8xf16>, tensor<2x2x1x8xf16>) -> tensor<2x2x1x8xf16>
+
+  %masked_f32 = tosa.cast %masked : (tensor<2x2x1x8xf16>) -> tensor<2x2x1x8xf32>
+  %max = tosa.reduce_max %masked_f32 {axis = 3 : i32} : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x1xf32>
+  %max_broadcast = tosa.mul %max, %softmax_ones, %shift : (tensor<2x2x1x1xf32>, tensor<2x2x1x8xf32>, tensor<1xi8>) -> tensor<2x2x1x8xf32>
+  %normalized = tosa.sub %masked_f32, %max_broadcast : (tensor<2x2x1x8xf32>, tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf32>
+  %exp = tosa.exp %normalized : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf32>
+  %sum = tosa.reduce_sum %exp {axis = 3 : i32} : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x1xf32>
+  %sum_broadcast = tosa.mul %sum, %softmax_ones, %shift : (tensor<2x2x1x1xf32>, tensor<2x2x1x8xf32>, tensor<1xi8>) -> tensor<2x2x1x8xf32>
+  %reciprocal = tosa.reciprocal %sum_broadcast : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf32>
+  %softmax = tosa.mul %exp, %reciprocal, %shift : (tensor<2x2x1x8xf32>, tensor<2x2x1x8xf32>, tensor<1xi8>) -> tensor<2x2x1x8xf32>
+  %softmax_f16 = tosa.cast %softmax : (tensor<2x2x1x8xf32>) -> tensor<2x2x1x8xf16>
+  %softmax_collapsed = tensor.collapse_shape %softmax_f16 [[0, 1], [2], [3]] : tensor<2x2x1x8xf16> into tensor<4x1x8xf16>
+  %value = tensor.expand_shape %arg3 [[0, 1, 2, 3]] output_shape [2, 2, 8, 2] : tensor<64xf16> into tensor<2x2x8x2xf16>
+  %value_collapsed = tensor.collapse_shape %value [[0, 1], [2], [3]] : tensor<2x2x8x2xf16> into tensor<4x8x2xf16>
+  %output = tosa.matmul %softmax_collapsed, %value_collapsed, %zero, %zero : (tensor<4x1x8xf16>, tensor<4x8x2xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<4x1x2xf16>
+  %result = tensor.collapse_shape %output [[0, 1, 2]] : tensor<4x1x2xf16> into tensor<8xf16>
+  return %result : tensor<8xf16>
+}
+
diff --git a/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir b/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir
index 1fd16f9b75da..ec08a9435674 100644
--- a/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir
+++ b/mlir/test/fusion/nightly-misc-e2e/mixr-attention/f16/mixr-attention-sliding-window-kvcache-prefix-causal.mlir
@@ -22,7 +22,8 @@ module {
     %3 = migraphx.literal(dense<0xFC00> : tensor<1xf16>) : <1xf16, 1>
     %4 = migraphx.literal(dense<5.000000e-01> : tensor<1xf16>) : <1xf16, 1>
     %sliding_offset = migraphx.literal(dense<-1> : tensor<1xsi32>) : <1xsi32, 1>
-    %fixed_seq_len = migraphx.literal(dense<2> : tensor<2x1xsi32>) : <2x1xsi32, 1x1>
+    // currentSeqLen is an inclusive key index, so the full-cache value is K-1.
+    %fixed_seq_len = migraphx.literal(dense<1> : tensor<2x1xsi32>) : <2x1xsi32, 1x1>
     %seq_len = migraphx.clip %arg2, %fixed_seq_len, %fixed_seq_len : <2x1xsi32, 1x1>, <2x1xsi32, 1x1>, <2x1xsi32, 1x1> -> <2x1xsi32, 1x1>
     %5 = migraphx.reshape %arg0 {dims = [2, 6, 1, 2, 2]} : <2x6x2x2xf16, 24x4x2x1> -> <2x6x1x2x2xf16, 24x4x4x2x1>
     %6 = migraphx.multibroadcast %5 {out_dyn_dims = [], out_lens = [2, 6, 2, 2, 2]} : <2x6x1x2x2xf16, 24x4x4x2x1> -> <2x6x2x2x2xf16, 24x4x0x2x1>