Enable swizzling SMEM for transposed dot operand #474
Conversation
htyu
changed the title
| @@ -131,6 +131,8 @@ compared to 1*64 when the hasLeadingOffset is false. | |||
|
|
|||
| if (mfmaEnc) { | |||
| int kDimNum = dotOpEnc.getOpIdx() == 0 ? 1 : 0; | |||
| if (needTrans) | |||
There was a problem hiding this comment.
@htyu Is this for tt.trans?
In other cases, we don't set the needTrans field when creating a sharedLayout.
There was a problem hiding this comment.
Yes, it's for tt.trans like below. When creating a shared encoding, it looks ahead to see if there's a transpose prior to the dot.
qk = tl.dot(q, tl.trans(k), allow_tf32=ALLOW_TF32) * alpha
There was a problem hiding this comment.
I learned that tl.trans is a very tricky op.
Let me clarify. The 10% improvement is due to the enablement of swizzling?
There was a problem hiding this comment.
Yes, this change speeds up our model that has above trans and dot operation by 10%.
There was a problem hiding this comment.
BTW, tl.trans is used quite often in attention kernel of a transformer model, iiuc.
There was a problem hiding this comment.
With this change, do we need a change related to lines at: https://github.com/htyu/triton/blob/ab3aafb4a9158a60a4c3085a13b605f70488d6ff/include/triton/Dialect/TritonGPU/IR/TritonGPUAttrDefs.td#L166-L167
We have scenario of non-square tile now.
There was a problem hiding this comment.
no they should be good.
The changes here is only used to enable swizzling for operands of dot.
getMDim and getNDim are shapes of the result of dot.
There was a problem hiding this comment.
Never mind, this is related to mfma layout
Transposed operand will be accessed in an opposite order from the original operand. Enabling swizzling seems to help performance. I'm seeing 10% performance improvement for our internal model. This is a backport of ROCm#474.
Transposed operand will be accessed in an opposite order from the original operand. Enabling swizzling seems to help performance. I'm seeing 10% performance improvement for our internal model.