add per_token_quant_bf16_int8 kernel #939
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 @theNiemand, 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 introduces a significant performance optimization by adding a new CUDA kernel for per-token quantization from BFloat16 to INT8. This change aims to reduce memory footprint and improve computational efficiency for deep learning workloads, particularly in large language models. The new kernel is fully integrated into the existing Python API and comes with dedicated tests to validate its accuracy and performance.
Highlights
- New CUDA Kernel for Quantization: Introduced a new
per_token_quant_bf16_int8CUDA kernel, designed to efficiently quantize BFloat16 tensors to INT8 format on a per-token basis. This includes specialized implementations for common tensor dimensions (e.g., 16, 32, 64, 512, 1024, 3200, 4096, 12800) and vectorized processing for optimal performance. - Python Bindings and API Integration: The newly added CUDA kernel is exposed to Python via
pybind11inops_bindings.cppand integrated into thelightllm_kernel.opsmodule, providing a seamless Python API for users to leverage this quantization functionality. - Utility Function for INT8 Conversion: A new
float_to_int8_rndevice utility function has been added, utilizing CUDA intrinsics for accurate rounding to nearest and saturating conversion from FP32 to INT8, crucial for the quantization process. - Comprehensive Testing: A new test file
int8_quant_test.pyhas been added, including accuracy and performance benchmarks for theper_token_quant_bf16_int8kernel against a reference implementation, ensuring the correctness and efficiency of the new feature.
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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.
| 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 issue 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 is currently in preview and 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 to provide feedback.
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 new CUDA kernel for per-token bfloat16 to int8 quantization. A critical bug was identified in the generic kernel's quantization logic. Additionally, suggestions were made to improve code clarity by removing an unused kernel parameter.
| workspace1[i] = local_bf16; | ||
|
|
||
| fp32_t tmp = cvt_bf16_f32(local_bf16); | ||
| local_max = fmaxf(local_max, tmp); |
There was a problem hiding this comment.
The maximum value is calculated without considering the absolute values. For symmetric quantization, find the maximum of the absolute values to ensure correct scaling for tensors with negative values.
local_max = fmaxf(local_max, fabsf(tmp));
| Tensor contiguous_input = input.is_contiguous() ? input : input.contiguous(); | ||
| Tensor contiguous_scales = scales.is_contiguous() ? scales : scales.contiguous(); | ||
|
|
||
| const int64_t M = input.size(0); |
There was a problem hiding this comment.
The M variable is passed to the CUDA kernels but is unused. Remove the M parameter from the kernel definitions and call sites to avoid confusion.
No description provided.