misc: fix vector size calculation for fp4 #1702
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Summary of Changes
Hello @yzh119, I'm Gemini Code Assist1! 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 resolves a critical bug in how vector sizes are calculated for 128-bit alignment, particularly impacting the fp4 data type. By introducing a specialized utility function, get_vec_size_128b, the system now correctly determines the number of elements per 128-bit vector, ensuring accurate memory access patterns and preventing potential issues in CUDA kernels that rely on this calculation. This change improves the robustness and correctness of operations involving mixed-precision data types.
Highlights
- Corrected Vector Size Calculation: Introduced a new utility function,
get_vec_size_128b, to accurately determine the number of elements that fit into a 128-bit vector for various data types. - FP4 Data Type Support: Specifically addressed an issue where the
sizeofoperator incorrectly calculated the vector size for__nv_fp4_e2m1(4-bit floating point) due to C++'s byte-addressable memory model, ensuring proper alignment for this sub-byte type. - Widespread Integration: Updated multiple CUDA kernels and utility functions across
activation,cp_async, andnormmodules to utilize the newget_vec_size_128bfunction, replacing the previous16 / sizeof(T)calculation.
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
-
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. ↩
There was a problem hiding this comment.
Code Review
This pull request introduces a utility function get_vec_size_128b to correctly calculate vector sizes for 128-bit memory operations, specifically to handle the fp4 data type which has a sub-byte size. The change is propagated to all relevant parts of the codebase, replacing the previous hardcoded calculation. The approach is sound. I have a couple of suggestions for the new utility function in include/flashinfer/utils.cuh to improve its robustness and maintainability.
| template <typename T> | ||
| __host__ __device__ __forceinline__ constexpr size_t get_vec_size_128b() { | ||
| #if CUDA_VERSION >= 12080 | ||
| if constexpr (std::is_same_v<T, __nv_fp4_e2m1>) { | ||
| return 32; // 128 bits / 4 bits per element = 32 elements | ||
| } else { | ||
| return 16 / sizeof(T); | ||
| } | ||
| #else | ||
| return 16 / sizeof(T); | ||
| #endif | ||
| } |
There was a problem hiding this comment.
This function can be simplified to improve readability and reduce code duplication. Additionally, it currently only handles the __nv_fp4_e2m1 type for FP4. The other FP4 type, __nv_fp4_e2m0, should also be handled to make the utility more robust.
Here's a suggested refactoring that addresses both points:
template <typename T>
__host__ __device__ __forceinline__ constexpr size_t get_vec_size_128b() {
#if CUDA_VERSION >= 12080
if constexpr (std::is_same_v<T, __nv_fp4_e2m1> || std::is_same_v<T, __nv_fp4_e2m0>) {
return 32; // 128 bits / 4 bits per element = 32 elements
}
#endif
return 16 / sizeof(T);
}The if constexpr with a return inside makes the else branch unnecessary. The default case return 16 / sizeof(T); can be shared for both CUDA versions and for types other than FP4.
template <typename T>
__host__ __device__ __forceinline__ constexpr size_t get_vec_size_128b() {
#if CUDA_VERSION >= 12080
if constexpr (std::is_same_v<T, __nv_fp4_e2m1> || std::is_same_v<T, __nv_fp4_e2m0>) {
return 32; // 128 bits / 4 bits per element = 32 elements
}
#endif
return 16 / sizeof(T);
}
| */ | ||
| #ifndef FLASHINFER_UTILS_CUH_ | ||
| #define FLASHINFER_UTILS_CUH_ | ||
| #include <cuda.h> |
There was a problem hiding this comment.
The <cuda.h> header seems redundant here. The CUDA_VERSION macro, which is used for conditional compilation, is also available through <cuda_runtime.h>, which is already included. Removing this unnecessary include helps in keeping file dependencies clean and minimal.
📌 Description
We used to rely on
16 / sizeof(T)to calculate the vector size of a given data type to compose 128b, but it will fail for fp4 because sizeof(__nv_fp4_e2m1) is 1 (there is no native subbyte design in C++).This PR fixes the issue by adding a function
get_vec_size_128bwhich will return 32 for fp4.🔍 Related Issues
🚀 Pull Request Checklist
Thank you for contributing to FlashInfer! Before we review your pull request, please make sure the following items are complete.
✅ Pre-commit Checks
pre-commitby runningpip install pre-commit(or used your preferred method).pre-commit install.pre-commit run --all-filesand fixed any reported issues.🧪 Tests
unittest, etc.).Reviewer Notes