GPUAI-1250 - Flash Attention v2.04 module rotary cannot be used code fixed

[xiaoxiangAMD]

background
as we have installed the flash-attention on the machine or a docker image, we found the module rotary_emb could not be used

flash-attention install description
please refer to https://github.com/ROCm/flash-attention

rotary_emb install description
1 first clone the repository:https://github.com/ROCm/flash-attention.git
2 enter the directory: /csrc/rotary/
3 run the command on a terminal or a shell :python setup.py install (python3 is needed)

command for the test
run the import code after you have run python:from flash_attn.layers.rotary import apply_rotary_emb_func

result expected

in the directory /tests,we get the test script for the rotary ,the result is as follows::

the success

[sabreshao]

Please give more test result beside importing.

[xiaoxiangAMD]

the test result for rotary module is given, the test source is in the file test_rotary.py