[attention] Extend attention to fuse transpose

[antiagainst]

No description provided.

[antiagainst]

Update 5/22: patch iree-org/iree#17408 out; needing review.
Update 5/23: working on decomposition and tiling. Patch out today or so.

[Groverkss]

Plan to finish it this week (Before Jun 7):

4 Jun: Land online attention (iree-org/iree#17536)
5 Jun: Create transform script using online_attention for MFMA
6 Jun: Add indexing_maps to attention op
7 Jun: Fusions

[raikonenfnu]

Hey guys, quick update

1. Indexing attention ([LinalgExt] Adding IndexingMaps to linalg_ext.attentionOp iree-org/iree#17864) has landed
2. CastTypeOFitMMA support for TD pipeline ([LLVMGPU] Support CastTypeToFitMMA on TransformDialect script.  iree-org/iree#17884) is up
3. transfer_write distribution for non contiguous indexing map ([LLVMGPU][VectorDist] Enable support to distribute vector.transfer_write with non-contiguous dims iree-org/iree#17895) is up

Once 2. and 3. and iree-org/iree@d2ca774 is landed on main, we should be able to handle/compile fused attn-transpose.

[antiagainst]

Awesome. All 3 pull requests are in. Can you send out the last piece?

[raikonenfnu]

Awesome. All 3 pull requests are in. Can you send out the last piece?

Hey Lei, I think @MaheshRavishankar is en route to pushing that one in! :)

[MaheshRavishankar]

I can send it in early next week.

[raikonenfnu]

I also pushed up/updated the spec mlir to find k2 correctly (link). I tested compiling on the fusion-preprocessing test MLIR (here) and was able to get a vmfb out.

The gist above is slightly different from the test in where we make the scale constant here. It fails on vector distribution if scale is not constant.

compile command:

~/nod/iree-build-notrace/tools/iree-compile constant_transpose_fusion.mlir --iree-hal-target-backends=rocm --iree-rocm-target-chip=gfx942 --iree-global-opt-propagate-transposes=true --iree-opt-outer-dim-concat=true --iree-opt-const-eval=false --iree-opt-data-tiling=false --iree-rocm-waves-per-eu=2 --iree-vm-target-truncate-unsupported-floats --iree-codegen-llvmgpu-use-vector-distribution --iree-codegen-gpu-native-math-precision=true --iree-flow-enable-aggressive-fusion -o attention.vmfb --iree-codegen-transform-dialect-library=attention_and_matmul_spec.mlir

[raikonenfnu]

FYI I also tested the attention-transpose-fusion vmfb numerics on normal random numbers (0.0, 1.0) against torch, seems like we have good numerics there :)

Starting IR, compile command, data generator can all be found in https://gist.github.com/raikonenfnu/973b4d91e4378702ce4b4496d732cb57

Needed to update the shape from the original fusion-preprocessing test a little bit since the fastest dim for Q,K,V needs to be the same to run on pytorch.