Skip to content

[AIROCMLIR-499] Add support for sliding window masking for attention #2240

New issue

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

Sign up for GitHub

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

Already on GitHub? Sign in to your account

Open
justinrosner wants to merge 12 commits into
base: develop
Choose a base branch
from
+825 −67
Open

[AIROCMLIR-499] Add support for sliding window masking for attention #2240

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
12 commits
Select commit Hold shift + click to select a range
65c86ef
Initial TosaToRock changes for sliding window attention
justinrosner Feb 12, 2026
ff90aeb
Add GridwiseGemmToBlockwise masking support
justinrosner Feb 12, 2026
2eda9a7
Small bug fixes
justinrosner Feb 12, 2026
2dd857f
Small refactor
justinrosner Feb 17, 2026
04c0048
More refactoring
justinrosner Feb 17, 2026
c76bd00
Clang-format
justinrosner Feb 17, 2026
3c5b79d
Add LIT tests
justinrosner Feb 17, 2026
dd4d525
Attend to review comments
justinrosner Feb 18, 2026
52a7058
Add rocmlir-gen and PR LIT tests
justinrosner Feb 18, 2026
fa10ee5
Add E2E tests
justinrosner Feb 18, 2026
873edc2
Clang-format
justinrosner Feb 18, 2026
10cd27e
Attend to more review comments
justinrosner Feb 24, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
12 changes: 10 additions & 2 deletions mlir/include/mlir/Dialect/Rock/IR/RockOps.td
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -220,7 +220,8 @@ def Rock_AttentionOp
Optional<TensorOrMemRefOf<[F32, F16, BF16]>>:$lse, I32Attr:$numHeadsQ,
I32Attr:$numHeadsKV, UnitAttr:$qTransposed, UnitAttr:$kTransposed,
UnitAttr:$vTransposed, UnitAttr:$oTransposed, UnitAttr:$causal,
I32Attr:$splitKV, OptionalAttr<Rock_GemmFeaturesAttr>:$features,
I32Attr:$splitKV, OptionalAttr<I32Attr>:$slidingWindowSize,
OptionalAttr<Rock_GemmFeaturesAttr>:$features,
StoreMethodAttr:$storeMethod, OptionalAttr<TypeAttr>:$softmaxType,
OptionalAttr<RockTuningParamAttrInterface>:$params0,
OptionalAttr<RockTuningParamAttrInterface>:$params1,
Expand Down Expand Up @@ -253,6 +254,11 @@ 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.

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.
Flash decoding multiplies the number of blocks by splitKV. Note that "lse" has to be non-null for splitKV > 1.
Expand All @@ -278,6 +284,7 @@ def Rock_AttentionOp
` ` `qk` `=` (`tr` $qTransposed^)? $queries `*` (`tr` $kTransposed^)? $keys `:` type($queries) `,` type($keys) `\n`
(`currentSeqLen` `=` `(` $currentSeqLen^ `:` type($currentSeqLen) `)` `\n`)?
(`prefixOffset` `=` `(` $prefixOffset^ `:` type($prefixOffset) `)` `\n`)?
(`slidingWindowSize` `=` $slidingWindowSize^ `\n`)?
(`causal` `\n` $causal^)?
(`lse` `=` $lse^ `:` type($lse) `\n`)?
(`qk` `=` `elementwise` (`otherIns` `(` $preSoftmaxElemWiseInputs^ `:` type($preSoftmaxElemWiseInputs) `)`)? $preSoftmaxBody^ `\n`)?
Expand Down Expand Up @@ -583,7 +590,8 @@ def Rock_GridwiseAttentionAccelOp
Optional<MemRefRankOf<[I32], [1]>>:$prefixOffset,
MemRefRankOf<[F32, F16, BF16], [3]>:$out,
Optional<MemRefRankOf<[F32, F16, BF16], [2]>>:$lse, UnitAttr:$causal,
I32Attr:$splitKV, OptionalAttr<Rock_GemmFeaturesAttr>:$features,
I32Attr:$splitKV, OptionalAttr<I32Attr>:$slidingWindowSize,
OptionalAttr<Rock_GemmFeaturesAttr>:$features,
StoreMethodAttr:$storeMethod, I32Attr:$blockSize, I32Attr:$gridSize,
UnitAttr:$disableQBypassLDS, OptionalAttr<IndexAttr>:$prePadG0M,
OptionalAttr<IndexAttr>:$prePadG0N,
Expand Down
Loading