[Model]fix_rwkv_cuda

[AswarS]

This PR adds external CUDA operators for RWKV and integrates them into the existing RWKV model implementation in FlagScale.

New Files:
Added 4 new CUDA-related source files implementing custom RWKV operators:
rwkv_cuda_op.cu
rwkv_cuda_kernel.cu
rwkv_cuda_wrapper.cpp
rwkv_cuda_utils.h
These files define and register RWKV-specific CUDA kernels to accelerate core computations and improve numerical stability.

Modified Files:
rwkv_model.py
Fixed the original rwkv_cuda computation logic.
Replaced internal computation paths with calls to the newly added external CUDA operators.

[gemini-code-assist]

Summary of Changes

Hello @AswarS, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the RWKV model's performance within the FlagScale framework by introducing and integrating external CUDA and HIP operators. It replaces the existing computation logic with highly optimized GPU kernels, ensuring faster execution and improved numerical stability for both forward and backward passes, particularly leveraging bfloat16 precision.

Highlights

• New CUDA Operators: Dedicated CUDA kernels (wkv7_cuda.cu, wkv7_hip.hip) have been introduced for the RWKV model's core operations, significantly enhancing computational performance and numerical stability.
• PyTorch Integration: Custom PyTorch operators (wkv7_op.cpp, wkv7_op.hip) are now implemented to seamlessly integrate these new CUDA and HIP kernels into the existing Python model, allowing for direct calls from the PyTorch framework.
• Model Logic Refactoring: The rwkv_model.py file has been updated to utilize these newly added external CUDA operators, replacing previous internal computation logic and removing a less efficient CPU/PyTorch fallback implementation.
• Bfloat16 Support and Assertions: The new WindBackstepping autograd function now includes explicit assertions to ensure inputs are of bfloat16 data type and are contiguous, promoting efficient and stable GPU computations.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces custom CUDA and HIP operators to accelerate RWKV model computations, integrating them into the existing Python model implementation. The changes are a good step towards improving performance. However, I've identified a few critical issues that need to be addressed. Firstly, there's a bug in the HIP operator registration that would cause it to fail on ROCm platforms. Secondly, the removal of a fallback implementation for the custom kernel will lead to a crash if the kernel is not enabled or fails to compile. Lastly, there are significant inconsistencies in variable naming across the Python, C++, and CUDA/HIP code, which harms readability and maintainability. I've provided detailed comments and suggestions to resolve these issues.

[CLAassistant]

Thank you for your submission! We really appreciate it. Like many open source projects, we ask that you sign our Contributor License Agreement before we can accept your contribution.
_{You have signed the CLA already but the status is still pending? Let us recheck it.}