diff --git a/mlir/include/mlir/Dialect/Rock/IR/RockOps.td b/mlir/include/mlir/Dialect/Rock/IR/RockOps.td index f0319adf16f8..1a92aae6ce3f 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. + + 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 09f3448ae452..a434a1a8d099 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" @@ -53,6 +54,7 @@ #include "llvm/Support/Debug.h" #include "llvm/Support/LogicalResult.h" #include "llvm/Support/raw_ostream.h" +#include #include #include #include @@ -1749,7 +1751,7 @@ struct AttentionMatcherValues { Value currentSeqLen; bool isCausal; Value prefixOffset; - std::optional slidingWindowSize; + std::optional slidingWindowSize; std::optional seqLenClipMin; std::optional seqLenClipMax; Type softmaxType; @@ -1964,27 +1966,71 @@ struct AttentionRewritePattern : public OpRewritePattern { return maybeSelect; } - // Helper to verify a value is i32 and traces back to a block argument - bool isI32BlockArgument(Value val, - const DenseSet &seqLenSkip) const { + // Resolve an i32 sequence-length value to its block argument. A non-trivial + // transpose or reshape is discarded only when it preserves the group-value + // interpretation expected by attention lowering. + FailureOr + resolveSeqLenBlockArgument(Value val, const DenseSet &seqLenSkip, + int64_t expectedNumGroups) const { auto shape = dyn_cast(val.getType()); if (!shape || !shape.getElementType().isInteger(32)) - return false; + return failure(); - FailureOr maybeBlockArg = getValueSkipping(val, seqLenSkip); - return succeeded(maybeBlockArg) && - isa(maybeBlockArg.value()); + while (Operation *definingOp = val.getDefiningOp()) { + if (!seqLenSkip.contains(definingOp->getName().getStringRef())) + break; + if (auto transpose = dyn_cast(definingOp)) { + ArrayRef inputShape = + cast(transpose.getInput1().getType()).getShape(); + for (auto [outputDim, inputDim] : + llvm::enumerate(transpose.getPerms())) { + if (inputDim != static_cast(outputDim) && + (inputShape[outputDim] != 1 || inputShape[inputDim] != 1)) + return failure(); + } + val = transpose.getInput1(); + } else if (isa(definingOp)) { + FailureOr maybeBroadcast = mulBroadcast(val); + if (failed(maybeBroadcast)) + return failure(); + val = *maybeBroadcast; + } else if (auto expand = dyn_cast(definingOp)) { + ArrayRef inputShape = expand.getSrcType().getShape(); + ArrayRef 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(definingOp)) { + val = collapse.getSrc(); + } else { + return failure(); + } + } + if (!isa(val)) + return failure(); + return val; } // 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 &seqLenSkip) const { - FailureOr resolvedA = getValueSkipping(a, seqLenSkip); - FailureOr resolvedB = getValueSkipping(b, seqLenSkip); + const DenseSet &seqLenSkip, + int64_t expectedNumGroups) const { + FailureOr resolvedA = + resolveSeqLenBlockArgument(a, seqLenSkip, expectedNumGroups); + FailureOr resolvedB = + resolveSeqLenBlockArgument(b, seqLenSkip, expectedNumGroups); return succeeded(resolvedA) && succeeded(resolvedB) && - isa(resolvedA.value()) && - resolvedA.value() == resolvedB.value(); + *resolvedA == *resolvedB; } // Helper function to detect select-based causal mask pattern: @@ -2111,22 +2157,31 @@ struct AttentionRewritePattern : public OpRewritePattern { Value inputToContinue; // The value to continue pattern matching with Value seqLen; // The sequence length Value prefixOffset; // The prefix offset value - std::optional slidingWindowSize; // The sliding window size + std::optional slidingWindowSize; // The sliding window size // Clip bounds detected on currentSeqLen during KV-cache pattern matching. std::optional seqLenClipMin; std::optional 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 slidingWindowClipMin; std::optional slidingWindowClipMax; }; // Helper to try detecting prefix causal pattern: add(row_indices, offset) // Returns the offset value if successful - FailureOr - tryPrefixCausalPattern(Value input, - const DenseSet &seqLenSkip) const { + FailureOr tryPrefixCausalPattern(Value input, + const DenseSet &seqLenSkip, + int64_t expectedNumGroups) const { DenseSet expandAndCollapse{ tensor::CollapseShapeOp::getOperationName(), tensor::ExpandShapeOp::getOperationName()}; @@ -2163,16 +2218,11 @@ struct AttentionRewritePattern : public OpRewritePattern { maybeOffset = offset; FailureOr 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; } @@ -2183,18 +2233,92 @@ struct AttentionRewritePattern : public OpRewritePattern { std::optional clipMax; }; - // Helper to try detecting KV-cache pattern. - // Also detects an optional clip (min(max(x, lo), hi)) on currentSeqLen. - FailureOr - tryKVCachePattern(Value input, const DenseSet &seqLenSkip) const { + struct ClipResult { + Value input; + std::optional clipMin; + std::optional clipMax; + }; + + // Peel all multiply-by-one operations used to broadcast a scalar-like value. + Value peelBroadcasts(Value input) const { + while (true) { + FailureOr maybeNonOne = mulBroadcast(input); + if (failed(maybeNonOne)) + return input; + input = *maybeNonOne; + } + } + + // Detect optional max(input, clipMin) and min(input, clipMax) bounds, + // allowing the constant to appear on either side of each commutative + // operation. + FailureOr tryClipPattern(Value input) const { DenseSet expandAndCollapse{ tensor::CollapseShapeOp::getOperationName(), tensor::ExpandShapeOp::getOperationName()}; - DenseSet expandCollapseMinMax{ - tensor::CollapseShapeOp::getOperationName(), - tensor::ExpandShapeOp::getOperationName(), - tosa::MaximumOp::getOperationName(), - tosa::MinimumOp::getOperationName()}; + + auto extractI32Constant = [&](Value value) -> std::optional { + 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(); + }; + + Value unclippedInput = input; + std::optional clipMin; + std::optional clipMax; + + auto maybeMin = + getDefiningOpSkipping(input, expandAndCollapse); + 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(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(); + } + } + + if (!clipMin && !clipMax) + return failure(); + return ClipResult{unclippedInput, clipMin, clipMax}; + } + + bool hasValidSeqLenClipBounds(const ClipResult &clip, + int64_t maxSeqLen) const { + if (maxSeqLen <= 0) + return false; + auto isOutOfRange = [maxSeqLen](std::optional bound) { + return bound && (*bound < 0 || static_cast(*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 + tryKVCachePattern(Value input, const DenseSet &seqLenSkip, + int64_t maxSeqLen, int64_t expectedNumGroups) const { FailureOr maybeNonOne = mulBroadcast(input); if (failed(maybeNonOne)) return failure(); @@ -2209,98 +2333,45 @@ struct AttentionRewritePattern : public OpRewritePattern { !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. + seqLenCandidate = peelBroadcasts(seqLenCandidate); + + auto maybeClip = tryClipPattern(seqLenCandidate); if (succeeded(maybeClip)) { + // 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; result.clipMax = maybeClip->clipMax; } - // Skip through expand/collapse/min/max to reach the block argument - auto maybeCurrentSeqLen = - getValueSkipping(maybeNonOne.value(), expandCollapseMinMax); - 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 layout-neutral transforms to the block argument so + // addBroadcastForBlockArg can reconstruct the head broadcast. + FailureOr maybeCurrentSeqLen = resolveSeqLenBlockArgument( + seqLenCandidate, seqLenSkip, expectedNumGroups); + if (failed(maybeCurrentSeqLen)) 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 tryClipPattern(Value input) const { - DenseSet expandAndCollapse{ - tensor::CollapseShapeOp::getOperationName(), - tensor::ExpandShapeOp::getOperationName()}; - - // Helper to extract a splat i32 constant from a value - auto extractI32Constant = [&](Value val) -> std::optional { - 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(); - }; - - // Look for tosa.minimum (the outer clip op) - auto maybeMin = - getDefiningOpSkipping(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 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(maxCandidate, expandAndCollapse); - if (failed(maybeMax)) - return failure(); - auto maxOp = maybeMax.value(); - - // One input of maximum is a constant (clipMin) - std::optional 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; + 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. 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 clipMin; std::optional clipMax; }; @@ -2309,16 +2380,11 @@ struct AttentionRewritePattern : public OpRewritePattern { // greater(add(seqLen, negative_const_offset) * broadcast, col_indices) // Returns the window size and validated currentSeqLen operand if successful. FailureOr - trySlidingWindowPattern(Value input, - const DenseSet &seqLenSkip) const { + trySlidingWindowPattern(Value input, const DenseSet &seqLenSkip, + int64_t maxSeqLen, int64_t expectedNumGroups) const { DenseSet expandAndCollapse{ tensor::CollapseShapeOp::getOperationName(), tensor::ExpandShapeOp::getOperationName()}; - DenseSet expandCollapseMinMax{ - tensor::CollapseShapeOp::getOperationName(), - tensor::ExpandShapeOp::getOperationName(), - tosa::MaximumOp::getOperationName(), - tosa::MinimumOp::getOperationName()}; // Trace through broadcast multiplication (mul by 1) FailureOr maybeNonOne = mulBroadcast(input); @@ -2363,24 +2429,40 @@ struct AttentionRewritePattern : public OpRewritePattern { 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::max() || + windowSize > maxSeqLen) + return failure(); + + // 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 clipMin; std::optional clipMax; - auto maybeClip = tryClipPattern(seqLenOperand); + Value seqLenCandidate = peelBroadcasts(seqLenOperand); + + auto maybeClip = tryClipPattern(seqLenCandidate); if (succeeded(maybeClip)) { + // 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; clipMax = maybeClip->clipMax; } // 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 maybeSeqLen = - getValueSkipping(seqLenOperand, expandCollapseMinMax); - Value seqLen = succeeded(maybeSeqLen) ? maybeSeqLen.value() : seqLenOperand; - if (!isI32BlockArgument(seqLen, seqLenSkip)) + FailureOr maybeSeqLen = resolveSeqLenBlockArgument( + seqLenCandidate, seqLenSkip, expectedNumGroups); + if (failed(maybeSeqLen)) return failure(); - return SlidingWindowResult{maybeWindowSize.value(), seqLen, clipMin, - clipMax}; + return SlidingWindowResult{static_cast(windowSize), *maybeSeqLen, + clipMin, clipMax}; } /* @@ -2544,7 +2626,9 @@ struct AttentionRewritePattern : public OpRewritePattern { void analyzeSelectForSeqLenMask(tosa::SelectOp select, SeqLenMaskResult &result, const DenseSet &opsToSkip, - const DenseSet &seqLenSkip) const { + const DenseSet &seqLenSkip, + int64_t maxSeqLen, + int64_t expectedNumGroups) const { auto pred = select.getInput1(); auto maybeGreater = getDefiningOpSkipping(pred, opsToSkip); if (failed(maybeGreater)) @@ -2559,7 +2643,8 @@ struct AttentionRewritePattern : public OpRewritePattern { // 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; @@ -2570,7 +2655,8 @@ struct AttentionRewritePattern : public OpRewritePattern { // 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(); } @@ -2583,9 +2669,14 @@ struct AttentionRewritePattern : public OpRewritePattern { 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); + auto maybeSlidingWindow = trySlidingWindowPattern( + input1, seqLenSkip, maxSeqLen, expectedNumGroups); if (succeeded(maybeSlidingWindow)) { auto slidingWindow = maybeSlidingWindow.value(); result.slidingWindowSize = slidingWindow.windowSize; @@ -2598,7 +2689,9 @@ struct AttentionRewritePattern : public OpRewritePattern { } } - FailureOr getSeqLenMask(Value softmaxInput) const { + FailureOr getSeqLenMask(Value softmaxInput, + int64_t maxSeqLen, + int64_t expectedNumGroups) const { auto maybeSelect = getSelectWithNegInf(softmaxInput); if (failed(maybeSelect)) return failure(); @@ -2610,13 +2703,13 @@ struct AttentionRewritePattern : public OpRewritePattern { 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 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, @@ -2624,14 +2717,17 @@ struct AttentionRewritePattern : public OpRewritePattern { nullptr, std::nullopt, std::nullopt}; // Analyze the first (outer) select - analyzeSelectForSeqLenMask(select, currentResult, opsToSkip, seqLenSkip); + analyzeSelectForSeqLenMask(select, currentResult, opsToSkip, seqLenSkip, + maxSeqLen, expectedNumGroups); // 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. + // 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); @@ -2646,7 +2742,7 @@ struct AttentionRewritePattern : public OpRewritePattern { auto chainedSelect = maybeChainedSelect.value(); int before = recognizedMaskCount(currentResult); analyzeSelectForSeqLenMask(chainedSelect, currentResult, opsToSkip, - seqLenSkip); + seqLenSkip, maxSeqLen, expectedNumGroups); // Leave an unrecognized or duplicate mask in the elementwise region. if (recognizedMaskCount(currentResult) == before) break; @@ -2657,20 +2753,24 @@ struct AttentionRewritePattern : public OpRewritePattern { // 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, + expectedNumGroups)) + return failure(); + // 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(); } // We need at least one pattern to be detected @@ -3267,9 +3367,30 @@ struct AttentionRewritePattern : public OpRewritePattern { // 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 slidingWindowSize; + std::optional slidingWindowSize; std::optional 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. 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 softmaxInputFinder; + softmaxInputFinder.visit(softmaxInput); + FailureOr maybeSourceMatMul = + softmaxInputFinder.getFirstGemmBasedOp(); + if (failed(maybeSourceMatMul)) { + LLVM_DEBUG( + llvm::dbgs() + << "first matmul not found before sequence-length mask analysis\n"); + return failure(); + } + ArrayRef keyShape = + cast(maybeSourceMatMul->getB().getType()).getShape(); + int64_t expectedNumGroups = keyShape.front(); + int64_t maxSeqLen = keyShape.back(); + + auto maybeSeqLenMask = + getSeqLenMask(softmaxInput, maxSeqLen, expectedNumGroups); if (succeeded(maybeSeqLenMask)) { auto result = maybeSeqLenMask.value(); kvCacheInput = result.inputToContinue; @@ -3311,6 +3432,12 @@ struct AttentionRewritePattern : public OpRewritePattern { 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 matC = maybeFirstMatMul.value().getOutput(); ArrayRef shapeC = matC.getType().getShape(); @@ -3449,10 +3576,10 @@ struct AttentionRewritePattern : public OpRewritePattern { 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(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 61ec7ebe34b3..838a8428d369 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> @@ -413,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> @@ -445,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} : (tensor<1x2x8x2xf16>) -> tensor<1x2x2x8xf16> %collapsed = tensor.collapse_shape %extracted_slice [[0, 1], [2], [3]] : tensor<1x2x1x2xf16> into tensor<2x1x2xf16> @@ -453,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> @@ -463,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> @@ -489,3 +490,157 @@ 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} : (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} : (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> +} + +// 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} : (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} : (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/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} : (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} : (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} : (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 f0a8320536d4..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 @@ -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> @@ -208,6 +213,367 @@ 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} : (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} : (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 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} : (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} : (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. +// 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} : (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} : (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} : (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} : (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. 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-NOT: 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} : (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> + %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> + %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> + %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} : (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 3176e7d5adb6..afac2eadfb6d 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 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" + # 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..2c8f7c6ade83 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 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" + # 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..647094b3cf29 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 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" + # 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..ad5b733c0219 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 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" + # 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/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> 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..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 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 - 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 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 sequenceLengthQ( "seq_len_q", llvm::cl::desc("sequence length of Q in attention()"), @@ -704,10 +704,13 @@ static llvm::cl::opt splitKV( static llvm::cl::opt 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 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