[sgl-kernel] Opt per_token_quant_fp8 with warp reduce

[yuan-luo]

Motivation

Optimize per_token_quant_fp8 kernel with warp reduce in large batch. Obtained 5-7% speedup.

Main:
$python ./sgl-kernel/benchmark/bench_per_token_quant_fp8.py
INFO 07-21 20:31:35 [__init__.py:244] Automatically detected platform cuda.
✅ All implementations match
per-token-dynamic-quant-fp8-performance:
    batch_size  seq_len         VLLM   SGL Kernel
0         16.0     64.0    28.511999    26.784001
1         16.0    128.0    46.239998    43.296002
2         16.0    256.0    84.384002    75.871997
3         16.0    512.0   156.544000   136.639997
4         16.0   1024.0   299.919993   257.216007
5         16.0   2048.0   587.743998   498.463988
6         16.0   4096.0  1164.703965   985.888004
7         32.0     64.0    46.080001    42.112000
8         32.0    128.0    84.416002    75.776003
9         32.0    256.0   156.703994   135.775998
10        32.0    512.0   300.224006   257.824004
11        32.0   1024.0   588.000000   498.400003
12        32.0   2048.0  1165.423989   985.519975
13        32.0   4096.0  2315.263987  1953.951955
14        64.0     64.0    84.416002    75.648002
15        64.0    128.0   156.448007   136.767998
16        64.0    256.0   300.511986   256.736010
17        64.0    512.0   587.423980   499.743998
18        64.0   1024.0  1165.279984   985.583991
19        64.0   2048.0  2314.559937  1953.855991
20        64.0   4096.0  4617.599964  3888.479948
21       128.0     64.0   156.448007   136.608005
22       128.0    128.0   300.543994   257.120013
23       128.0    256.0   587.776005   499.648005
24       128.0    512.0  1165.791988   985.775977
25       128.0   1024.0  2317.280054  1954.208016
26       128.0   2048.0  4613.759995  3888.384104
27       128.0   4096.0  9216.303825  7765.295982

This PR:
$python ./sgl-kernel/benchmark/bench_per_token_quant_fp8.py
INFO 07-21 20:23:15 [__init__.py:244] Automatically detected platform cuda.
✅ All implementations match
per-token-dynamic-quant-fp8-performance:
    batch_size  seq_len         VLLM   SGL Kernel
0         16.0     64.0    27.712001    26.496001
1         16.0    128.0    45.728002    42.688001
2         16.0    256.0    84.144000    77.344000
3         16.0    512.0   156.192005   143.776000
4         16.0   1024.0   299.584001   259.840012
5         16.0   2048.0   587.328017   493.088007
6         16.0   4096.0  1164.127946   951.968014
7         32.0     64.0    45.791999    41.312002
8         32.0    128.0    84.128000    76.288000
9         32.0    256.0   156.287998   143.808007
10        32.0    512.0   299.775988   260.127991
11        32.0   1024.0   587.375998   493.984014
12        32.0   2048.0  1165.312052   952.032030
13        32.0   4096.0  2316.672087  1857.887983
14        64.0     64.0    84.224001    76.127999
15        64.0    128.0   156.223997   142.976001
16        64.0    256.0   299.872011   260.031998
17        64.0    512.0   587.584019   493.903995
18        64.0   1024.0  1166.432023   951.968014
19        64.0   2048.0  2313.472033  1858.207941
20        64.0   4096.0  4616.223812  3668.463945
21       128.0     64.0   156.159997   143.647999
22       128.0    128.0   300.191998   259.871989
23       128.0    256.0   587.135971   494.031996
24       128.0    512.0  1165.375948   952.208012
25       128.0   1024.0  2313.695908  1858.335972
26       128.0   2048.0  4615.392208  3667.359948
27       128.0   4096.0  9214.896202  7294.496059

Modifications

Checklist

• Format your code according to the Code Formatting with Pre-Commit.
• Add unit tests as outlined in the Running Unit Tests.
• Update documentation / docstrings / example tutorials as needed, according to Writing Documentation.
• Provide throughput / latency benchmark results and accuracy evaluation results as needed, according to Benchmark and Profiling and Accuracy Results.
• For reviewers: If you haven't made any contributions to this PR and are only assisting with merging the main branch, please remove yourself as a co-author when merging the PR.
• Please feel free to join our Slack channel at https://slack.sglang.ai to discuss your PR.

[gemini-code-assist]

Summary of Changes

Hello @yuan-luo, 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 focuses on a significant performance optimization for the per_token_quant_fp8 CUDA kernel, which is critical for efficient FP8 quantization in the SGLang framework. By refactoring the kernel to utilize advanced CUDA features like warp-level parallelism and the cute tensor library, it achieves notable speed improvements across various batch and sequence lengths.

Highlights

• Performance Optimization: I've re-implemented the per_token_quant_fp8_kernel to leverage warp-level reductions and cute tensors for improved efficiency, resulting in up to 10% speedup as demonstrated by the provided benchmarks.
• Kernel Architecture Refinement: The kernel now transitions from block-level reduction (cub::block::block_reduce) to warp-level reduction (warpReduceMax) for calculating the maximum value per token, enabling more granular and efficient parallelization.
• CUDA cute Tensor Integration: I've adopted the cute tensor library for managing global memory access within the kernel, providing a more expressive and potentially optimized way to handle tensor operations.
• Optimized Kernel Launch Configuration: The kernel launch parameters have been adjusted to align with the new warp-centric design, processing multiple tokens per CTA (specifically, 8 tokens per 256-thread CTA).

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

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 refactors the per_token_quant_fp8 CUDA kernel to improve performance by processing multiple tokens per thread block, with each token handled by a single warp. This is a solid optimization strategy. The review identified a critical race condition in the use of shared memory for the scaling factor, which would lead to incorrect quantization results. A medium-severity issue with an unnecessary const_cast was also found.

[BBuf]

Based on the benchmark results, when batch_size and seq_length are small, the warp reduce indirectly reduces the number of blocks. This leads to an inability to fully utilize all SMs, resulting in performance degradation.

I recommend that you should determine whether warp_reduce optimization is necessary based on the number of SMs and num_tokens, rather than directly replacing the existing kernel implementation. Overall, it's crucial to ensure that the kernel's performance is superior in all scenarios.

[BBuf]

This optimization might be better addressed after the merger of https://github.com/sgl-project/sglang/pull/7604; otherwise, it could introduce significant conflicts.

[yuan-luo]

This optimization might be better addressed after the merger of https://github.com/sgl-project/sglang/pull/7604; otherwise, it could introduce significant conflicts.

@BBuf , [https://github.com//pull/7604] is for per_token_group_quant_int8/fp8, not related with this PR which is per_token_quant_fp8.

[yuan-luo]

Based on the benchmark results, when batch_size and seq_length are small, the warp reduce indirectly reduces the number of blocks. This leads to an inability to fully utilize all SMs, resulting in performance degradation.

I recommend that you should determine whether warp_reduce optimization is necessary based on the number of SMs and num_tokens, rather than directly replacing the existing kernel implementation. Overall, it's crucial to ensure that the kernel's performance is superior in all scenarios.

Refactored.
Now the performance improved in both small batch and large batch. Keep legacy function for small batch, introduce warp reduce for large batch.

$python ./sgl-kernel/benchmark/bench_per_token_quant_fp8.py
INFO 07-21 20:23:15 [__init__.py:244] Automatically detected platform cuda.
✅ All implementations match
per-token-dynamic-quant-fp8-performance:
    batch_size  seq_len         VLLM   SGL Kernel
0         16.0     64.0    27.712001    26.496001
1         16.0    128.0    45.728002    42.688001
2         16.0    256.0    84.144000    77.344000
3         16.0    512.0   156.192005   143.776000
4         16.0   1024.0   299.584001   259.840012
5         16.0   2048.0   587.328017   493.088007
6         16.0   4096.0  1164.127946   951.968014
7         32.0     64.0    45.791999    41.312002
8         32.0    128.0    84.128000    76.288000
9         32.0    256.0   156.287998   143.808007
10        32.0    512.0   299.775988   260.127991
11        32.0   1024.0   587.375998   493.984014
12        32.0   2048.0  1165.312052   952.032030
13        32.0   4096.0  2316.672087  1857.887983
14        64.0     64.0    84.224001    76.127999
15        64.0    128.0   156.223997   142.976001
16        64.0    256.0   299.872011   260.031998
17        64.0    512.0   587.584019   493.903995
18        64.0   1024.0  1166.432023   951.968014
19        64.0   2048.0  2313.472033  1858.207941
20        64.0   4096.0  4616.223812  3668.463945
21       128.0     64.0   156.159997   143.647999
22       128.0    128.0   300.191998   259.871989
23       128.0    256.0   587.135971   494.031996
24       128.0    512.0  1165.375948   952.208012
25       128.0   1024.0  2313.695908  1858.335972
26       128.0   2048.0  4615.392208  3667.359948
27       128.0   4096.0  9214.896202  7294.496059

[BBuf]

LGTM now. @ispobock Can you have a look too? thanks!