[TTIR] Add SDPA fusing pattern #5999
Open
+1,565
−86
Conversation
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
jserbedzijaTT
requested review from
a team,
arminaleTT,
ctodTT,
dlokeTT,
jgrimTT,
jnie-TT,
mrakitaTT,
odjuricicTT,
sgholamiTT and
tapspatel
as code owners
jserbedzijaTT
requested review from
AleksKnezevic,
acolicTT,
azecevicTT,
bmalesevicTT,
kmabeeTT,
mtopalovicTT,
nvukobratTT,
pilkicTT,
rpavlovicTT,
sdjordjevicTT,
staylorTT,
svuckovicTT and
vwellsTT
1 task
jserbedzijaTT
commented
Nov 23, 2025
Comment on lines
+2073
to
+2122
| auto queryType = mlir::cast<RankedTensorType>(op.getQuery().getType()); | ||
|
|
||
| // If sequence length of query is 1, we should use the decode op | ||
| if (queryType.getDimSize(2) == 1) { | ||
| // Permute query: [B, H, 1, D] -> [1, B, H, D] | ||
| Value permutedQuery = ttir_to_ttnn::utils::generatePermute( | ||
| mlir::cast<TypedValue<mlir::RankedTensorType>>(adaptor.getQuery()), | ||
| {2, 0, 1, 3}, rewriter, op.getLoc()); | ||
|
|
||
| // Broadcast mask head dimension if needed for decode op | ||
| Value attentionMask = adaptor.getAttentionMask(); | ||
| if (attentionMask) { | ||
| auto maskType = mlir::cast<RankedTensorType>(attentionMask.getType()); | ||
| int64_t numHeads = queryType.getDimSize(1); | ||
|
|
||
| if (numHeads > 1) { | ||
| SmallVector<int64_t> broadcastShape(maskType.getShape().begin(), | ||
| maskType.getShape().end()); | ||
| broadcastShape[2] = numHeads; | ||
|
|
||
| auto broadcastType = ttnn::utils::RankedTensorTypeFactory::create( | ||
| maskType, broadcastShape); | ||
| auto broadcastDims = ttmlir::utils::getBroadcastDimensions<int64_t>( | ||
| maskType.getShape(), broadcastShape); | ||
| auto shapeAttr = | ||
| ttnn::ShapeAttr::get(rewriter.getContext(), broadcastDims); | ||
| attentionMask = rewriter.create<ttnn::RepeatOp>( | ||
| op.getLoc(), broadcastType, attentionMask, shapeAttr); | ||
| } | ||
| } | ||
|
|
||
| auto decodeOp = rewriter.create<ttnn::ScaledDotProductAttentionDecodeOp>( | ||
| op.getLoc(), permutedQuery.getType(), permutedQuery, adaptor.getKey(), | ||
| adaptor.getValue(), op.getIsCausal(), attentionMask, | ||
| /*cur_pos_tensor=*/Value(), /*attention_sink=*/Value(), | ||
| adaptor.getScaleAttr(), /*memory_config=*/nullptr, | ||
| /*program_config=*/nullptr); | ||
|
|
||
| // Permute result back: [1, B, H, D] -> [B, H, 1, D] | ||
| rewriter.replaceOp( | ||
| op, ttir_to_ttnn::utils::generatePermute( | ||
| decodeOp.getResult(), {1, 2, 0, 3}, rewriter, op.getLoc())); | ||
| } else { | ||
| rewriter.replaceOpWithNewOp<ttnn::ScaledDotProductAttentionOp>( | ||
| op, this->getTypeConverter()->convertType(op.getType()), | ||
| adaptor.getQuery(), adaptor.getKey(), adaptor.getValue(), | ||
| adaptor.getAttentionMask(), op.getIsCausal(), adaptor.getScaleAttr(), | ||
| adaptor.getSlidingWindowSizeAttr(), | ||
| /*memory_config=*/nullptr); | ||
| } |
Contributor
Author
There was a problem hiding this comment.
As mentioned in the description, this is probably not the best place to put the logic of deciding whether to create ttnn.scaled_dot_product_attention or ttnn.scaled_dot_product_attention_decode operation so if anyone have a better idea feel free to comment.
Codecov Report❌ Patch coverage is Additional details and impacted files@@ Coverage Diff @@
## main #5999 +/- ##
==========================================
- Coverage 69.28% 69.21% -0.08%
==========================================
Files 333 335 +2
Lines 50899 51193 +294
==========================================
+ Hits 35264 35431 +167
- Misses 15635 15762 +127 ☔ View full report in Codecov by Sentry. |
Contributor
|
@jserbedzijaTT Great! Have you tested this on llama e2e (details)? If so, could you share the resulting IR? |
sgholamiTT
approved these changes
Nov 25, 2025
| OptionalAttr<F32Attr>:$scale, | ||
| OptionalAttr<TTNN_MemoryConfigAttr>:$memory_config); | ||
| OptionalAttr<TTNN_MemoryConfigAttr>:$memory_config, | ||
| OptionalAttr<TTNN_SDPAProgramConfigAttr>:$program_config); |
Contributor
There was a problem hiding this comment.
I believe we should pass this information to optimizer (it's missing here).
In general, the optimizer code should be similar to the runtime code.
We also need whether this logic needs to change or not with this modification.
jserbedzijaTT
force-pushed
the
jserbedzija/add_sdpa_ttir_fusing_pattern
branch
2 times, most recently
from
798847d to
04ab34f
Compare
jserbedzijaTT
force-pushed
the
jserbedzija/add_sdpa_ttir_fusing_pattern
branch
from
04ab34f to
73aaa10
Compare
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.
This suggestion is invalid because no changes were made to the code.
Suggestions cannot be applied while the pull request is closed.
Suggestions cannot be applied while viewing a subset of changes.
Only one suggestion per line can be applied in a batch.
Add this suggestion to a batch that can be applied as a single commit.
Applying suggestions on deleted lines is not supported.
You must change the existing code in this line in order to create a valid suggestion.
Outdated suggestions cannot be applied.
This suggestion has been applied or marked resolved.
Suggestions cannot be applied from pending reviews.
Suggestions cannot be applied on multi-line comments.
Suggestions cannot be applied while the pull request is queued to merge.
Suggestion cannot be applied right now. Please check back later.
Ticket
Link to Github Issue
Problem description
Scaled Dot Product Attention (SDPA) is a fundamental building block in transformer architectures, computing
Currently, we lack a fusing pattern in TTIR to combine this chain of operations into a single
ttir.scaled_dot_product_attentionop.What's changed
SDPA Fusing Pattern:
TTIRFusingpass that identifies the canonical attention computation pattern:ttir.scaled_dot_product_attentionoperation[batch_size, 1, query_seq_len, key_seq_len]if neededLowering of
ttir.scaled_dot_product_attentionttir.scaled_dot_product_attentionin theTTIRToTTNNpass. It now lowers to eitherttnn.scaled_dot_product_attentionorttnn.scaled_dot_product_attention_decodedepending on thesequence_lengthdimension in thequerytensor (ifsequence_lengthis 1, it uses the decode operation).Workarounds:
ttnn.scaled_dot_product_attention_decodeop requiresprogram_configto succeed so we create a default one in the workaround.Testing:
Checklist