roll back fmha/common.py

[tenpercent]

Addressing

The current interface of MHA is that H has to always match between qkv. If you want to do GQA - e.g. one kv-head for every n q-heads, you have to send 5D inputs. (Thus we're forcing the user to be very explicit.) Do we really want to relax that rule in this PR?

And also merging test_mqa_forward into test_mqa_decoding as suggested in

To follow up on @bottler's question about @rocm_only - it'd be better for fmha.ck.FwOp to be covered by the generic test_forward and test_mqa_decoding. Then we don't need a separate test function (and eventually won't need @rocm_only after all such cases are refactored)

Not sure if this should be blocking the merge or can be done as a follow-up

[qianfengz]

@tenpercent

• The current ck.Fw.Op together with its underlying ck-tiled implementation is able to support mqa/gqa, even though the supported input tensors are 4-D
• But the ref_attention in test_mem_eff_attention.py is not able to handle mqa/gqa with 4-D inputs, so that is why ref_mqa_attention is added
• test_mqa_forward is added for explicitly verified those functions

[qianfengz]

Since we added too many scripts for this function, we can just remove them, and I will keep the scripts in private for testing/verification.