[kernel] add kernels for llama inference

[tiandiao123]

📌 Checklist before creating the PR

• I have created an issue for this PR for traceability
• The title follows the standard format: [doc/gemini/tensor/...]: A concise description
• I have added relevant tags if possible for us to better distinguish different PRs

🚨 Issue number

Link this PR to your issue with words like fixed to automatically close the linked issue upon merge

e.g. fixed #1234, closed #1234, resolved #1234

📝 What does this PR do?

This PR is used to add llama useful cuda kernels for inference.

💥 Checklist before requesting a review

• I have linked my PR to an issue (instruction)
• My issue clearly describes the problem/feature/proposal, with diagrams/charts/table/code if possible
• I have performed a self-review of my code
• I have added thorough tests.
• I have added docstrings for all the functions/methods I implemented

⭐️ Do you enjoy contributing to Colossal-AI?

• 🌝 Yes, I do.
• 🌚 No, I don't.

Tell us more if you don't enjoy contributing to Colossal-AI.

[kurisusnowdeng]

Any differece from the original codes? For example, I quickly checked pos_encoding & layernorm kernels but they seem exactly the same as the codes in vllm's repository.
If this is just integrating some external codes into our project, I recommend just to add the external tool as a dependency and import the necessary API for our needs.

[tiandiao123]

Any differece from the original codes? For example, I quickly checked pos_encoding & layernorm kernels but they seem exactly the same as the codes in vllm's repository. If this is just integrating some external codes into our project, I recommend just to add the external tool as a dependency and import the necessary API for our needs.

There is some minor changes. I only integrate useful parts which will be integrated into our llama-inference attention over here.

[github-actions]

The code coverage for the changed files is %.

Click me to view the complete report

Name                                               Stmts   Miss  Cover
----------------------------------------------------------------------
op_builder/__init__.py                                11      0   100%
op_builder/rmsnorm.py                                 21      9    57%
op_builder/rotary_embedding.py                        42     25    40%
setup.py                                              73     73     0%
tests/test_kernels/cuda/test_rmsnorm.py               42      4    90%
tests/test_kernels/cuda/test_rotary_embedding.py      75      4    95%
tests/test_kernels/triton/test_self_attention.py      74     56    24%
tests/test_kernels/triton/test_softmax.py             17      8    53%
----------------------------------------------------------------------
TOTAL                                                355    179    50%

[tiandiao123]

Adding llama shardformer demo, try to directly import third-party useful ops @kurisusnowdeng

[tiandiao123]

use this PR instead: #4485 @kurisusnowdeng

[tiandiao123]

temporarily closed this.