Skip to content

Fix attention mask traversal for three nested selects #2437

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
umangyadav merged 7 commits into from
Jul 31, 2026
+209 −26
Merged

Fix attention mask traversal for three nested selects #2437

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
868969d
Fix sliding-window sequence length mask detection
umangyadav Jul 23, 2026
ba15f65
Fix attention mask traversal for three nested selects
umangyadav Jul 23, 2026
381912b
[NFC] Make three-mask folding checks order-independent
umangyadav Jul 23, 2026
e41ac51
Merge branch 'develop' into users/umayadav/fix-attention-three-mask
umangyadav Jul 27, 2026
0b362d8
Merge branch 'develop' into users/umayadav/fix-attention-three-mask
umangyadav Jul 28, 2026
5c4f12b
[NFC] Remove duplicate sliding-window mask reconciliation
umangyadav Jul 29, 2026
a114b78
Merge branch 'develop' into users/umayadav/fix-attention-three-mask
umangyadav Jul 31, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
51 changes: 25 additions & 26 deletions mlir/lib/Conversion/TosaToRock/TosaToRock.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2626,32 +2626,31 @@ struct AttentionRewritePattern : public OpRewritePattern<tosa::MatMulOp> {
// Analyze the first (outer) select
analyzeSelectForSeqLenMask(select, currentResult, opsToSkip, seqLenSkip);

// Check if the inputToContinue (input3) is another chained select with
// -inf. This handles cases where multiple mask patterns (KVCache, prefix
// causal, sliding window) use separate selects.
bool haveSeqLen = currentResult.seqLen != nullptr;
bool havePrefixOffset = currentResult.prefixOffset != nullptr;
bool haveSlidingWindow = currentResult.slidingWindowSize.has_value();

// Try chaining if we found at least one pattern but not all
bool foundAny = haveSeqLen || havePrefixOffset || haveSlidingWindow;
bool foundAll = haveSeqLen && havePrefixOffset && haveSlidingWindow;
if (foundAny && !foundAll) {
auto maybeChainedSelect = getSelectWithNegInf(inputToContinue);
if (succeeded(maybeChainedSelect)) {
auto chainedSelect = maybeChainedSelect.value();
// Try to analyze the chained select for the missing pattern
analyzeSelectForSeqLenMask(chainedSelect, currentResult, opsToSkip,
seqLenSkip);
// Only update inputToContinue if we found a complementary pattern
bool foundComplementary =
(!haveSeqLen && currentResult.seqLen) ||
(!havePrefixOffset && currentResult.prefixOffset) ||
(!haveSlidingWindow && currentResult.slidingWindowSize.has_value());
if (foundComplementary) {
currentResult.inputToContinue = chainedSelect.getInput3();
}
}
// Iteratively peel chained select(mask, -inf, scores) ops to detect
// separately nested KV-cache, prefix-causal, and sliding-window masks.
// 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.
auto recognizedMaskCount = [](const SeqLenMaskResult &result) {
return (result.seqLen ? 1 : 0) + (result.prefixOffset ? 1 : 0) +
(result.slidingWindowSize.has_value() ? 1 : 0);
};
while (recognizedMaskCount(currentResult) > 0 &&
recognizedMaskCount(currentResult) < 3) {
auto maybeChainedSelect =
getSelectWithNegInf(currentResult.inputToContinue);
if (failed(maybeChainedSelect))
break;

auto chainedSelect = maybeChainedSelect.value();
int before = recognizedMaskCount(currentResult);
analyzeSelectForSeqLenMask(chainedSelect, currentResult, opsToSkip,
seqLenSkip);
// Leave an unrecognized or duplicate mask in the elementwise region.

@justinrosner justinrosner Jul 29, 2026

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Not necessarily urgent for this PR, but in rocmlirTriton (assuming we are going to port this change there) we may run into errors with the duplicate/unrecognized masks being left in the elementwise region. The regularize* passes cannot sink transforms through non-splat constants (I've seen errors with this before when we were first testing the FusionZoo kernels with the rocmlirTriton prototype)

if (recognizedMaskCount(currentResult) == before)
break;
currentResult.inputToContinue = chainedSelect.getInput3();
}

// Sliding-window masking is defined relative to currentSeqLen. Reconcile
Expand Down
94 changes: 94 additions & 0 deletions mlir/test/Conversion/TosaToRock/tosa-to-rock-attention-three-mask.mlir
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
// RUN: sed s/##TOKEN_ARCH##/%arch/g %s | rocmlir-opt --tosa-to-rock -verify-diagnostics | FileCheck %s

// Three nested select ops must all be folded into a single rock.attention.
// From inner to outer, the masks are prefix-causal, KV-cache, and sliding
// window. All four corresponding attention properties must be preserved.
// CHECK-LABEL: func @attention_three_mask
// CHECK: rock.attention
// CHECK-DAG: currentSeqLen = (%{{.*}}
// CHECK-DAG: prefixOffset = (%{{.*}}
// CHECK-DAG: causal
// CHECK-DAG: slidingWindowSize = 3
// CHECK: qk = elementwise {
// CHECK-NOT: tosa.select
// CHECK: rock.yield

module {
func.func @attention_three_mask(%arg0: tensor<1xi32>, %arg1: tensor<9216xf16>, %arg2: tensor<2048xf16>, %arg3: tensor<2048xf16>, %arg4: tensor<2xi32>) -> tensor<7168xf16> attributes {rock.kernel, rock.arch = "##TOKEN_ARCH##"} {
%0 = "tosa.const"() <{values = dense<[[0, 1, 2, 3, 4, 5, 6, 7]]> : tensor<1x8xi32>}> : () -> tensor<1x8xi32>
%1 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<2x14x4x8xf32>}> : () -> tensor<2x14x4x8xf32>
%2 = "tosa.const"() <{values = dense<0xFC00> : tensor<2x14x4x8xf16>}> : () -> tensor<2x14x4x8xf16>
%3 = "tosa.const"() <{values = dense<0.000000e+00> : tensor<1xf16>}> : () -> tensor<1xf16>
%4 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<2x2x7x64x8xf16>}> : () -> tensor<2x2x7x64x8xf16>
%5 = "tosa.const"() <{values = dense<1.000000e+00> : tensor<2x2x7x8x64xf16>}> : () -> tensor<2x2x7x8x64xf16>
%6 = "tosa.const"() <{values = dense<1> : tensor<2x14x4x8xi8>}> : () -> tensor<2x14x4x8xi8>
%7 = "tosa.const"() <{values = dense<1.250000e-01> : tensor<2x14x4x8xf16>}> : () -> tensor<2x14x4x8xf16>
%8 = "tosa.const"() <{values = dense<0> : tensor<1xi8>}> : () -> tensor<1xi8>
%9 = "tosa.const"() <{values = dense<1> : tensor<4x8xi32>}> : () -> tensor<4x8xi32>
%10 = "tosa.const"() <{values = dense<1> : tensor<4x1xi32>}> : () -> tensor<4x1xi32>
%11 = "tosa.const"() <{values = dense<1> : tensor<2x8xi32>}> : () -> tensor<2x8xi32>
%12 = "tosa.const"() <{values = dense<[[0], [1], [2], [3]]> : tensor<4x1xi32>}> : () -> tensor<4x1xi32>
%swoff = "tosa.const"() <{values = dense<-3> : tensor<1x1xi32>}> : () -> tensor<1x1xi32>
%expanded = tensor.expand_shape %arg1 [[0, 1, 2, 3]] output_shape [2, 4, 18, 64] : tensor<9216xf16> into tensor<2x4x18x64xf16>
%13 = tosa.transpose %expanded {perms = array<i32: 0, 2, 1, 3>} : (tensor<2x4x18x64xf16>) -> tensor<2x18x4x64xf16>
%expanded_0 = tensor.expand_shape %arg0 [[0, 1]] output_shape [1, 1] : tensor<1xi32> into tensor<1x1xi32>
%14 = tosa.mul %expanded_0, %10, %8 : (tensor<1x1xi32>, tensor<4x1xi32>, tensor<1xi8>) -> tensor<4x1xi32>
%15 = tosa.add %14, %12 : (tensor<4x1xi32>, tensor<4x1xi32>) -> tensor<4x1xi32>
%16 = tosa.mul %15, %9, %8 : (tensor<4x1xi32>, tensor<4x8xi32>, tensor<1xi8>) -> tensor<4x8xi32>
%17 = tosa.mul %0, %9, %8 : (tensor<1x8xi32>, tensor<4x8xi32>, tensor<1xi8>) -> tensor<4x8xi32>
%18 = tosa.greater %17, %16 : (tensor<4x8xi32>, tensor<4x8xi32>) -> tensor<4x8xi1>
%19 = tosa.cast %18 : (tensor<4x8xi1>) -> tensor<4x8xi32>
%20 = tosa.cast %19 : (tensor<4x8xi32>) -> tensor<4x8xi8>
%expanded_1 = tensor.expand_shape %20 [[0, 1, 2], [3]] output_shape [1, 1, 4, 8] : tensor<4x8xi8> into tensor<1x1x4x8xi8>
%21 = tosa.mul %expanded_1, %6, %8 : (tensor<1x1x4x8xi8>, tensor<2x14x4x8xi8>, tensor<1xi8>) -> tensor<2x14x4x8xi8>
%22 = tosa.mul %0, %11, %8 : (tensor<1x8xi32>, tensor<2x8xi32>, tensor<1xi8>) -> tensor<2x8xi32>
%expanded_2 = tensor.expand_shape %arg4 [[0, 1]] output_shape [2, 1] : tensor<2xi32> into tensor<2x1xi32>
%23 = tosa.mul %expanded_2, %11, %8 : (tensor<2x1xi32>, tensor<2x8xi32>, tensor<1xi8>) -> tensor<2x8xi32>
%24 = tosa.greater %22, %23 : (tensor<2x8xi32>, tensor<2x8xi32>) -> tensor<2x8xi1>
%25 = tosa.cast %24 : (tensor<2x8xi1>) -> tensor<2x8xi32>
%26 = tosa.cast %25 : (tensor<2x8xi32>) -> tensor<2x8xi8>
%expanded_3 = tensor.expand_shape %26 [[0, 1, 2], [3]] output_shape [2, 1, 1, 8] : tensor<2x8xi8> into tensor<2x1x1x8xi8>
%27 = tosa.mul %expanded_3, %6, %8 : (tensor<2x1x1x8xi8>, tensor<2x14x4x8xi8>, tensor<1xi8>) -> tensor<2x14x4x8xi8>
%sw0 = tosa.add %expanded_2, %swoff : (tensor<2x1xi32>, tensor<1x1xi32>) -> tensor<2x1xi32>
%sw1 = tosa.mul %sw0, %11, %8 : (tensor<2x1xi32>, tensor<2x8xi32>, tensor<1xi8>) -> tensor<2x8xi32>
%sw2 = tosa.greater %sw1, %22 : (tensor<2x8xi32>, tensor<2x8xi32>) -> tensor<2x8xi1>
%sw3 = tosa.cast %sw2 : (tensor<2x8xi1>) -> tensor<2x8xi32>
%sw4 = tosa.cast %sw3 : (tensor<2x8xi32>) -> tensor<2x8xi8>
%expanded_sw = tensor.expand_shape %sw4 [[0, 1, 2], [3]] output_shape [2, 1, 1, 8] : tensor<2x8xi8> into tensor<2x1x1x8xi8>
%sw5 = tosa.mul %expanded_sw, %6, %8 : (tensor<2x1x1x8xi8>, tensor<2x14x4x8xi8>, tensor<1xi8>) -> tensor<2x14x4x8xi8>
%extracted_slice = tensor.extract_slice %13[0, 0, 0, 0] [2, 14, 4, 64] [1, 1, 1, 1] : tensor<2x18x4x64xf16> to tensor<2x14x4x64xf16>
%expanded_4 = tensor.expand_shape %arg2 [[0, 1, 2, 3, 4]] output_shape [2, 2, 1, 8, 64] : tensor<2048xf16> into tensor<2x2x1x8x64xf16>
%28 = tosa.mul %expanded_4, %5, %8 : (tensor<2x2x1x8x64xf16>, tensor<2x2x7x8x64xf16>, tensor<1xi8>) -> tensor<2x2x7x8x64xf16>
%expanded_5 = tensor.expand_shape %arg3 [[0, 1, 2, 3, 4]] output_shape [2, 2, 1, 8, 64] : tensor<2048xf16> into tensor<2x2x1x8x64xf16>
%29 = tosa.transpose %expanded_5 {perms = array<i32: 0, 1, 2, 4, 3>} : (tensor<2x2x1x8x64xf16>) -> tensor<2x2x1x64x8xf16>
%30 = tosa.mul %29, %4, %8 : (tensor<2x2x1x64x8xf16>, tensor<2x2x7x64x8xf16>, tensor<1xi8>) -> tensor<2x2x7x64x8xf16>
%collapsed = tensor.collapse_shape %extracted_slice [[0, 1], [2], [3]] : tensor<2x14x4x64xf16> into tensor<28x4x64xf16>
%collapsed_6 = tensor.collapse_shape %30 [[0, 1, 2], [3], [4]] : tensor<2x2x7x64x8xf16> into tensor<28x64x8xf16>
%31 = tosa.matmul %collapsed, %collapsed_6, %3, %3 {acc_type = f32} : (tensor<28x4x64xf16>, tensor<28x64x8xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<28x4x8xf16>
%expanded_7 = tensor.expand_shape %31 [[0, 1], [2], [3]] output_shape [2, 14, 4, 8] : tensor<28x4x8xf16> into tensor<2x14x4x8xf16>
%32 = tosa.mul %expanded_7, %7, %8 : (tensor<2x14x4x8xf16>, tensor<2x14x4x8xf16>, tensor<1xi8>) -> tensor<2x14x4x8xf16>
%33 = tosa.cast %21 : (tensor<2x14x4x8xi8>) -> tensor<2x14x4x8xi1>
%34 = tosa.select %33, %2, %32 : (tensor<2x14x4x8xi1>, tensor<2x14x4x8xf16>, tensor<2x14x4x8xf16>) -> tensor<2x14x4x8xf16>
%35 = tosa.cast %27 : (tensor<2x14x4x8xi8>) -> tensor<2x14x4x8xi1>
%36 = tosa.select %35, %2, %34 : (tensor<2x14x4x8xi1>, tensor<2x14x4x8xf16>, tensor<2x14x4x8xf16>) -> tensor<2x14x4x8xf16>
%sw6 = tosa.cast %sw5 : (tensor<2x14x4x8xi8>) -> tensor<2x14x4x8xi1>
%sw7 = tosa.select %sw6, %2, %36 : (tensor<2x14x4x8xi1>, tensor<2x14x4x8xf16>, tensor<2x14x4x8xf16>) -> tensor<2x14x4x8xf16>
%37 = tosa.cast %sw7 : (tensor<2x14x4x8xf16>) -> tensor<2x14x4x8xf32>
%38 = tosa.reduce_max %37 {axis = 3 : i32} : (tensor<2x14x4x8xf32>) -> tensor<2x14x4x1xf32>
%39 = tosa.mul %38, %1, %8 : (tensor<2x14x4x1xf32>, tensor<2x14x4x8xf32>, tensor<1xi8>) -> tensor<2x14x4x8xf32>
%40 = tosa.sub %37, %39 : (tensor<2x14x4x8xf32>, tensor<2x14x4x8xf32>) -> tensor<2x14x4x8xf32>
%41 = tosa.exp %40 : (tensor<2x14x4x8xf32>) -> tensor<2x14x4x8xf32>
%42 = tosa.reduce_sum %41 {axis = 3 : i32} : (tensor<2x14x4x8xf32>) -> tensor<2x14x4x1xf32>
%43 = tosa.mul %42, %1, %8 : (tensor<2x14x4x1xf32>, tensor<2x14x4x8xf32>, tensor<1xi8>) -> tensor<2x14x4x8xf32>
%44 = tosa.reciprocal %43 : (tensor<2x14x4x8xf32>) -> tensor<2x14x4x8xf32>
%45 = tosa.mul %41, %44, %8 : (tensor<2x14x4x8xf32>, tensor<2x14x4x8xf32>, tensor<1xi8>) -> tensor<2x14x4x8xf32>
%46 = tosa.cast %45 : (tensor<2x14x4x8xf32>) -> tensor<2x14x4x8xf16>
%collapsed_8 = tensor.collapse_shape %46 [[0, 1], [2], [3]] : tensor<2x14x4x8xf16> into tensor<28x4x8xf16>
%collapsed_9 = tensor.collapse_shape %28 [[0, 1, 2], [3], [4]] : tensor<2x2x7x8x64xf16> into tensor<28x8x64xf16>
%47 = tosa.matmul %collapsed_8, %collapsed_9, %3, %3 {acc_type = f32} : (tensor<28x4x8xf16>, tensor<28x8x64xf16>, tensor<1xf16>, tensor<1xf16>) -> tensor<28x4x64xf16>
%expanded_10 = tensor.expand_shape %47 [[0, 1], [2], [3]] output_shape [2, 14, 4, 64] : tensor<28x4x64xf16> into tensor<2x14x4x64xf16>
%48 = tosa.transpose %expanded_10 {perms = array<i32: 0, 2, 1, 3>} : (tensor<2x14x4x64xf16>) -> tensor<2x4x14x64xf16>
%collapsed_11 = tensor.collapse_shape %48 [[0, 1, 2, 3]] : tensor<2x4x14x64xf16> into tensor<7168xf16>
return %collapsed_11 : tensor<7168xf16>
}
}
Loading
Loading