[Re-land][jit kernel] Support per_token_group_quant_8bit jit kernel by yuan-luo · Pull Request #19140 · sgl-project/sglang

yuan-luo · 2026-02-22T04:32:22Z

Motivation

#18905 was reverted due to CI failure, but this CI failure was not introduced by #18905. After confirmed with @hnyls2002 , reland this jit kernel.

Adopted the fix in #19118 and added @mmangkad as co-author.

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Format your code according to the Format code with pre-commit.
Add unit tests according to the Run and add unit tests.
Update documentation according to Write documentations.
Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
Follow the SGLang code style guidance.

Review Process

Ping Merge Oncalls to start the PR flow. See the PR Merge Process.
Get approvals from CODEOWNERS and other reviewers.
Trigger CI tests with comments or contact authorized users to do so.
- /tag-run-ci-label, /rerun-failed-ci, /tag-and-rerun-ci
After green CI and required approvals, ask Merge Oncalls to merge.

gemini-code-assist · 2026-02-22T04:32:38Z

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 re-introducing and integrating a Just-In-Time (JIT) compiled kernel for per-token group 8-bit quantization. The primary goal is to enhance the efficiency of quantization operations by leveraging a dedicated C++/CUDA implementation. This involves adding the kernel's source, creating its Python interface, and ensuring its functional correctness and performance through new testing and benchmarking infrastructure.

Highlights

Re-land of JIT Kernel: This pull request re-lands the per_token_group_quant_8bit JIT kernel, which was previously reverted due to an unrelated CI failure. The kernel supports 8-bit quantization for per-token groups.
New C++/CUDA Kernel Implementation: A new C++/CUDA kernel (per_token_group_quant_8bit.cuh) has been added, implementing the core logic for per-token group 8-bit quantization, including support for column-major scales and UE8M0 format.
Python Wrapper and Integration: A Python wrapper (per_token_group_quant_8bit.py) has been introduced to expose the new JIT kernel to the Python codebase, and the fp8_kernel.py module has been updated to conditionally utilize this JIT kernel.
Benchmarking and Testing: New benchmark and unit test files have been added to validate the performance and correctness of the per_token_group_quant_8bit JIT kernel against an existing Triton implementation across various configurations.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Changelog

python/sglang/jit_kernel/benchmark/bench_per_token_group_quant_8bit.py
- Added a new benchmark script to compare the performance of the SGLang JIT kernel against a Triton implementation for per-token group 8-bit quantization.
python/sglang/jit_kernel/csrc/gemm/per_token_group_quant_8bit.cuh
- Added a new C++/CUDA kernel for per-token group 8-bit quantization, including logic for handling column-major scales and UE8M0 format.
python/sglang/jit_kernel/per_token_group_quant_8bit.py
- Added a Python wrapper for the new JIT kernel, allowing it to be called from Python with appropriate tensor and parameter handling.
python/sglang/jit_kernel/tests/test_per_token_group_quant_8bit.py
- Added comprehensive unit tests for the new JIT kernel, verifying its correctness against an existing Triton implementation across various configurations.
python/sglang/jit_kernel/utils.py
- Updated the CPP_DTYPE_MAP to include fp8_e4m3_t and int8_t for C++ type mapping.
- Removed an unnecessary blank line.
python/sglang/srt/layers/quantization/fp8_kernel.py
- Modified the sglang_per_token_group_quant_fp8 function to conditionally use the new JIT kernel (sgl_per_token_group_quant_8bit_jit) when enable_v2 is false, otherwise falling back to the existing sgl_per_token_group_quant_8bit (non-JIT) implementation.

Activity

The pull request is a re-land of [jit kernel] Support per_token_group_quant_8bit jit kernel #18905, which was previously reverted due to an unrelated CI failure.
The provided CI log indicates a failure in test_piecewise_cuda_graph_small_1_gpu.py, which appears to be an existing issue and not introduced by this PR's changes.

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.

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

The pull request re-lands the JIT kernel for per_token_group_quant_8bit. The implementation is consistent with the repository's JIT infrastructure and provides a fallback path for the pre-compiled extension. I have identified a few issues that need to be addressed: missing support for torch.float8_e4m3fnuz in the JIT utility mapping which will break AMD support, and several hardcoded values in the test script that prevent it from correctly verifying different quantization configurations and dtypes. Additionally, there are minor cleanup and optimization opportunities in the CUDA kernel.

gemini-code-assist · 2026-02-22T04:34:40Z

+    DST_DTYPE* __restrict__ output_q,
+    scale_packed_t_t<kScaleUE8M0>* __restrict__ output_s,
+    const int group_size,
+    const int num_groups,


The num_groups parameter is unused in the kernel (as indicated by the (void)num_groups; on line 50). It should be removed from the signature.

yuan-luo · 2026-02-22T04:41:57Z

/rerun-failed-ci

yuan-luo · 2026-02-22T04:42:06Z

/tag-and-rerun-ci

Co-authored-by: Mohammad Miadh Angkad <mangkad.bsdsba2027@aim.edu>

yuan-luo · 2026-02-22T05:27:54Z

/tag-and-rerun-ci

yuan-luo · 2026-02-22T07:32:01Z

/rerun-failed-ci

yuan-luo · 2026-02-22T13:16:17Z

/rerun-failed-ci

yuan-luo · 2026-02-22T15:01:27Z

/rerun-failed-ci

yuan-luo · 2026-02-23T00:19:31Z

/rerun-failed-ci

yuan-luo · 2026-02-23T01:42:40Z

/rerun-failed-ci

yuan-luo · 2026-02-23T06:19:41Z

/rerun-failed-ci

yuan-luo · 2026-02-24T03:05:08Z

/rerun-failed-ci

yuan-luo · 2026-02-24T11:31:06Z

/rerun-failed-ci

yuan-luo · 2026-02-25T11:12:56Z

/rerun-failed-ci

yuan-luo · 2026-02-25T15:21:40Z

/rerun-failed-ci

BBuf · 2026-02-26T01:59:53Z

-                masked_m,
-                enable_v2=enable_v2,
-            )
+            if enable_v2:


Todo: after remove per_token_group_quant_8bit kernels to jit_kernel, we should unify per_token_group_quant_8bit and per_token_group_quant_8bit_v2 to a single file.

Sure. I'll follow up.

…gl-project#19140) Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com> Co-authored-by: Mohammad Miadh Angkad <mangkad.bsdsba2027@aim.edu>

…kernel (#19140)" This reverts commit 7c9e8e2.

…gl-project#19140) Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com> Co-authored-by: Mohammad Miadh Angkad <mangkad.bsdsba2027@aim.edu>

luoyuan.luo added 4 commits February 22, 2026 12:27

Support per_token_group_quant_8bit jit kernel

f4b4cd5

Fix errors and refactor code

71cc487

Add benchmark test

0a595a8

Address review comments

7070576

yuan-luo requested review from AniZpZ, BBuf, DarkSharpness, Edwardf0t1, FlamingoPg and ch-wan as code owners February 22, 2026 04:32

github-actions Bot added the quant LLM Quantization label Feb 22, 2026

gemini-code-assist Bot reviewed Feb 22, 2026

View reviewed changes

github-actions Bot added the run-ci label Feb 22, 2026

yuan-luo mentioned this pull request Feb 22, 2026

Fix Dynamo tracing in per_token_group_quant_8bit #19118

Closed

5 tasks

yuan-luo marked this pull request as draft February 22, 2026 05:01

yuan-luo force-pushed the jit_per_token_group_quant branch from 73ea775 to 7070576 Compare February 22, 2026 05:05

Fix Dynamo tracing in per_token_group_quant_8bit

cead1a9

Co-authored-by: Mohammad Miadh Angkad <mangkad.bsdsba2027@aim.edu>

yuan-luo marked this pull request as ready for review February 22, 2026 05:27

yuan-luo requested a review from hnyls2002 February 22, 2026 05:38

BBuf approved these changes Feb 22, 2026

View reviewed changes

DarkSharpness approved these changes Feb 24, 2026

View reviewed changes

BBuf merged commit 7c9e8e2 into sgl-project:main Feb 26, 2026
945 of 1032 checks passed

BBuf mentioned this pull request Feb 26, 2026

[Feature] sgl-kernel wheel slimming plan tracking #17865

Closed

74 tasks

BBuf reviewed Feb 26, 2026

View reviewed changes

hnyls2002 pushed a commit that referenced this pull request Mar 9, 2026

Revert "[Re-land][jit kernel] Support per_token_group_quant_8bit jit …

2024c31

…kernel (#19140)" This reverts commit 7c9e8e2.

This was referenced Mar 9, 2026

Revert "[Re-land][jit kernel] Support per_token_group_quant_8bit jit kernel (#19140)" #20164

Closed

[jit kernel] per_token_group_quant_8bit breaks test_mimo_models bs=1 on H200 #20165

Closed

yuan-luo deleted the jit_per_token_group_quant branch March 9, 2026 15:07

Conversation

yuan-luo commented Feb 22, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Motivation

Modifications

Accuracy Tests

Benchmarking and Profiling

Checklist

Review Process

Uh oh!

gemini-code-assist Bot commented Feb 22, 2026

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

gemini-code-assist Bot Feb 22, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

yuan-luo commented Feb 22, 2026

Uh oh!

yuan-luo commented Feb 22, 2026

Uh oh!

yuan-luo commented Feb 22, 2026

Uh oh!

yuan-luo commented Feb 22, 2026

Uh oh!

yuan-luo commented Feb 22, 2026

Uh oh!

yuan-luo commented Feb 22, 2026

Uh oh!

yuan-luo commented Feb 23, 2026

Uh oh!

yuan-luo commented Feb 23, 2026

Uh oh!

yuan-luo commented Feb 23, 2026

Uh oh!

yuan-luo commented Feb 24, 2026

Uh oh!

yuan-luo commented Feb 24, 2026

Uh oh!

yuan-luo commented Feb 25, 2026

Uh oh!

yuan-luo commented Feb 25, 2026

Uh oh!

Uh oh!

BBuf Feb 26, 2026

Choose a reason for hiding this comment

Uh oh!

yuan-luo Feb 26, 2026

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

yuan-luo commented Feb 22, 2026 •

edited

Loading