[XLA:GPU] Do not create transpose ops with non-default layout in `D…
#430
Conversation
…otDecomposer`. The `DotDecomposer` pass runs ahead of layout assignment. Introducing non-default layouts at this stage causes complications for subsequent passes, in particular the `DotMerger` pass. It also prevents generating fat fused triton kernels with output_tiles [1, 1], which leads to register spilling. PiperOrigin-RevId: 824476578
i-chaochen
left a comment
i-chaochen
left a comment
There was a problem hiding this comment.
after we have this, we don't need to disable trion_softmax anymore? 3febc29#diff-778936bfc8e1f6ad8322e9795a405576e306dcd8f2479dc3d29b0b9a6f51a832R551
Just to be aware of that jax is already released 0.7.1 and they are targeting 0.8 now. Could you check this whether in 0.8?
trion_softmax flag is by default enabled in 0.7.1. In upstream, there is no such flag, triton_softmax is always enabled. |
|
Thanks! then please backport it to 0.8 as well. |
|
please use the PR #432. |
Do not create
transposeops with non-default layout in DotDecomposer.The
DotDecomposerpass runs ahead of layout assignment. Introducing non-default layouts at this stage causes complications for subsequent passes, in particular theDotMergerpass.It also prevents generating fat fused triton kernels with output_tiles [1, 1], which leads to register spilling.
related issue - https://github.com/ROCm/frameworks-internal/issues/13595
related PR - #428