Opinion on Refactoring Ulysses

Is your feature request related to a problem? Please describe.
Hi, I am suggesting my ideas on refactoring Ulysses.

1. post_all2all_func and .permute() seems unnecessary if we opt out for regular all_to_all function instead of all_to_all_single. After all_to_all, we can torch.cat along the gather_idx dimension to concatenate agnostically.
2. Double nested if loops due to batch_dim_idx makes code very hard to read. I'm thinking that splitting input tensor agnostic to batch_dim_idx and using all_to_all can also alleviate this?

Questions from Commit 17ed7c7

1. So interestingly, fusing QKV actually degrades performance if concat and slices are necessary?
2. Do you think you can elaborate on pipelining QKV comm and GEMM? Is this referring to doing GEMM with data that stays during all2all?

Appreciate your feedback!

Review commit 17ed7c7 and the implementations of post_all2all_func, permute(), all_to_all, and all_to_all_single. The issue asks for opinions about a possible refactor and QKV communication rather than defining a patch or completion criteria, so seek maintainer direction before coding.