[2/2] Optimizations and refactors about quant kernel#7601
[2/2] Optimizations and refactors about quant kernel#7601fzyzcjy wants to merge 587 commits intosgl-project:mainfrom
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 @fzyzcjy, 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 optimizations and refactors to the quantization kernel, specifically focusing on per-token-group quantization for FP8 and INT8 data types. The changes include performance improvements, code unification, and the addition of a new scaling flag, enhancing the flexibility and efficiency of the quantization process.
Highlights
- Optimization: Optimized the per-token-group quantization kernel for both FP8 and INT8 data types.
- Refactor: Refactored the code to unify INT8 and FP8 quantization logic where possible, improving code maintainability.
- New Feature: Added support for a new
scale_ue8m0flag to control the scaling behavior during quantization. - Configuration: Introduced
PER_TOKEN_GROUP_QUANT_8BIT_VALID_FLAGSto define valid configurations for quantization parameters.
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. ↩
| def per_token_group_quant_8bit( | ||
| x: torch.Tensor, | ||
| group_size: int, | ||
| dst_dtype: torch.dtype, | ||
| eps: float = 1e-10, | ||
| column_major_scales: bool = False, | ||
| scale_tma_aligned: bool = False, | ||
| scale_ue8m0: bool = False, | ||
| ): | ||
| assert ( | ||
| x.shape[-1] % group_size == 0 | ||
| ), "the last dimension of `x` cannot be divisible by `group_size`" | ||
| assert x.is_contiguous(), "`x` is not contiguous" | ||
| ) -> Tuple[torch.Tensor, torch.Tensor]: | ||
| from sglang.srt.layers.quantization.int8_kernel import per_token_group_quant_int8 | ||
|
|
||
| if dst_dtype == torch.int8: | ||
| assert not column_major_scales | ||
| assert not scale_tma_aligned | ||
| assert not scale_ue8m0 | ||
| return per_token_group_quant_int8( | ||
| x=x, | ||
| group_size=group_size, | ||
| eps=eps, | ||
| dtype=dst_dtype, | ||
| ) | ||
|
|
||
| x_q = torch.empty_like(x, device=x.device, dtype=fp8_dtype) | ||
| return per_token_group_quant_fp8( | ||
| x=x, | ||
| group_size=group_size, | ||
| eps=eps, | ||
| column_major_scales=column_major_scales, |
There was a problem hiding this comment.
The per_token_group_quant_8bit function is missing the scale_ue8m0 parameter in its signature, but it's called with scale_ue8m0 as a keyword argument in the tests. This will cause a TypeError at runtime. Additionally, scale_ue8m0 should be asserted False for int8 quantization, as it is not supported.
def per_token_group_quant_8bit(
x: torch.Tensor,
group_size: int,
dst_dtype: torch.dtype,
eps: float = 1e-10,
column_major_scales: bool = False,
scale_tma_aligned: bool = False,
scale_ue8m0: bool = False,
) -> Tuple[torch.Tensor, torch.Tensor]:
from sglang.srt.layers.quantization.int8_kernel import per_token_group_quant_int8
if dst_dtype == torch.int8:
assert not column_major_scales
assert not scale_tma_aligned
assert not scale_ue8m0, "scale_ue8m0 is not supported for int8 quantization"
return per_token_group_quant_int8(
x=x,
group_size=group_size,
eps=eps,
dtype=dst_dtype,
)
return per_token_group_quant_fp8(
x=x,
group_size=group_size,
eps=eps,
column_major_scales=column_major_scales,
scale_tma_aligned=scale_tma_aligned,
scale_ue8m0=scale_ue8m0,
)| scale_tma_aligned: bool, | ||
| scale_ue8m0: bool, | ||
| ): | ||
| if scale_ue8m0: |
There was a problem hiding this comment.
The group size is hardcoded to 128 here. Use the group_size parameter passed into this function to make it more generic.
x_s_mn, x_s_k = x_q_mn, x_q_k // group_size| dtype=torch.int, | ||
| ).transpose(0, 1)[:x_s_mn, :] | ||
| elif column_major_scales: | ||
| if scale_tma_aligned: |
There was a problem hiding this comment.
There is a TODO comment here to extract an align function. The align function is already imported from sglang.math_utils. Use it directly.
aligned_size = align(x_shape[-2], 4)
fzyzcjy
changed the title
fzyzcjy
force-pushed
the
feat/opt_quant_related
branch
2 times, most recently
from
dd7eb08 to
531478f
Compare
# Conflicts: # python/sglang/srt/bench_utils.py # python/sglang/srt/layers/quantization/fp8_kernel.py # sgl-kernel/benchmark/bench_per_token_group_quant_8bit.py # sgl-kernel/csrc/gemm/per_token_group_quant_8bit.cu # sgl-kernel/tests/test_per_token_group_quant_8bit.py
fzyzcjy
changed the title
|
FYI this is currently depending on releasing new sgl-kernel |
3 participants
EDIT: kernel part is in #9534
EDIT: the code is ready and has large speedup and is being reviewed, I mark as "draft" b/c do not want (frequent) code push to trigger all CIs.
EDIT: 9344 -> 9945, e2e 6.4%
Motivation
Modifications
Checklist