[JIT Kernel] Migrate per_token_quant_fp8 to JIT

[xingsy97]

Motivation

Migrate per_token_quant_fp8 kernel to JIT compilation (#17865).

Modifications

Under python/sglang/jit_kernel/:

• csrc/gemm/per_token_quant_fp8.cuh — CUDA kernel (ported from sgl-kernel/csrc/gemm/per_token_quant_fp8.cu)
• per_token_quant_fp8.py — Python wrapper
• tests/test_per_token_quant_fp8.py — Unit tests
• benchmark/bench_per_token_quant_fp8.py — Benchmark (JIT vs AOT vs PyTorch)

And

• python/sglang/srt/layers/quantization/fp8_kernel.py — Switch import to JIT

Accuracy Tests

Pass all tests defined in python/sglang/jit_kernel/tests/test_per_token_quant_fp8.py

Benchmarking and Profiling

GPU: NVIDIA GeForce RTX 5060 Laptop GPU
Driver: 580.126.20 | CUDA: 12.8 | PyTorch: 2.9.1+cu128
Script: python/sglang/jit_kernel/benchmark/bench_per_token_quant_fp8.py

per-token-quant-fp8-performance (unit: us):
   M × K (dtype)       JIT    AOT    PyTorch  AOT->JIT
   1 × 512, fp16       4.6    5.1    37.3     -9.8%
   1 × 4096, bf16      4.7    5.2    51.2     -9.6%
   16 × 2048, fp16     4.8    5.3    54.7     -9.4%
   64 × 4096, fp32     7.1    7.0    120.8    +1.4%
   128 × 7168, bf16    8.6    8.2    224.6    +4.9%
   256 × 7168, fp16    12.3   12.1   410.2    +1.7%

JIT and AOT performance are essentially the same, no regression observed. Both are significantly faster than PyTorch.

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.

[gemini-code-assist]

Summary of Changes

Hello, 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 significantly enhances the per_token_quant_fp8 operation by migrating its implementation to a JIT-compiled CUDA kernel. This change aims to improve the performance of FP8 quantization within the SGLang framework, as demonstrated by substantial speedups over the PyTorch reference implementation. The migration involves new kernel development, Python integration, and thorough testing and benchmarking to ensure both correctness and efficiency.

Highlights

• JIT Migration: Migrated the per_token_quant_fp8 operation to a JIT-compiled kernel for improved performance and integration within the SGLang framework.
• New CUDA Kernel: Introduced a new CUDA kernel (per_token_quant_fp8.cuh) that implements a two-pass FP8 quantization strategy: first finding the absolute maximum value via CTA reduction, then quantizing with vectorized loads/stores.
• Python Wrapper and Integration: Developed a new Python wrapper (per_token_quant_fp8.py) to facilitate JIT loading and execution of the new CUDA kernel, registered as a custom operation.
• Comprehensive Testing: Added extensive unit tests (test_per_token_quant_fp8.py) covering 138 parametrized cases to verify the accuracy of the JIT kernel against a PyTorch reference across various input shapes and data types.
• Performance Benchmarking: Included a new benchmark script (bench_per_token_quant_fp8.py) to compare the performance of the JIT kernel against AOT (Ahead-Of-Time) and PyTorch implementations, demonstrating significant speedups.
• Module Import Update: Updated the import path for per_token_quant_fp8 in fp8_kernel.py to correctly reference the newly integrated JIT kernel.

🧠 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_quant_fp8.py
  • Added a new benchmark script for per_token_quant_fp8 comparing JIT, AOT, and PyTorch implementations.
• python/sglang/jit_kernel/csrc/gemm/per_token_quant_fp8.cuh
  • Added a new CUDA kernel for per-token FP8 quantization, featuring a two-pass approach for absmax finding and vectorized quantization.
• python/sglang/jit_kernel/per_token_quant_fp8.py
  • Added a new Python wrapper for the JIT-compiled per_token_quant_fp8 CUDA kernel, including JIT loading and custom op registration.
• python/sglang/jit_kernel/tests/test_per_token_quant_fp8.py
  • Added new unit tests for per_token_quant_fp8 covering various tensor shapes, data types, and comparing against a PyTorch reference.
• python/sglang/srt/layers/quantization/fp8_kernel.py
  • Updated the import statement for sgl_per_token_quant_fp8 to point to the new JIT kernel implementation.

Activity

• Accuracy tests passed for 138 parametrized cases, covering various M, K, and dtype combinations, including [m,1]-shaped output_s tests.
• Benchmarking results show significant performance improvements for the JIT kernel compared to the PyTorch implementation across different input sizes and data types.
• The code has been formatted, unit tests added, and benchmark results provided as per the contribution guidelines.

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

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 successfully migrates the per_token_quant_fp8 function to a JIT-compiled CUDA kernel, which is a great step for kernel slimming. The new implementation is well-structured with a dedicated CUDA kernel, a Python wrapper, comprehensive unit tests, and benchmarks. The code is clean and the approach is solid. I have a couple of suggestions to enhance the benchmark accuracy and further optimize the CUDA kernel's performance.

[DarkSharpness]

Why use reinterpret_cast here? Why not AlignedVector?

[xingsy97]

Done. changed to AlignedVector

[DarkSharpness]

Can we use some utility function to avoid passing the arguments same many times? The only difference here is the kernel.

e.g. use a local lambda functon which returns a pointer to the real kernel

[xingsy97]

Done. Extracted common logic into a lambda.

[HydraQYH]

Why 32? Does this assume a Blackwell architecture?

[xingsy97]

32 bytes here matches the AOT original version. In per_token_quant_fp8.cu L18, the default template parameter is kVecSize = 16. For float16 (2 bytes per element), that's 16 * 2 = 32 bytes per vector load.

The formula 32 / sizeof(DType) generalizes this across dtypes (e.g., float32 gives 32 / 4 = 8, also matching the AOT use_vec8 path).

From my understanding, 32-byte vector loads work correctly on all architectures. On pre-Blackwell GPUs the compiler splits them into two 16-byte load instructions. If #19794 gets merged first, I can update this to use device::kMaxVecBytes introduced by it directly.

[DarkSharpness]

Let's wait for #19794. You may profile the performance first, just ensure changing from 32 -> 16 should not bring no performance regression.

[xingsy97]

posted benchmark result in #19702 (comment)

[HydraQYH]

32 bytes here matches the AOT original version. In per_token_quant_fp8.cu L18, the default template parameter is kVecSize = 16. For float16 (2 bytes per element), that's 16 * 2 = 32 bytes per vector load.

The formula 32 / sizeof(DType) generalizes this across dtypes (e.g., float32 gives 32 / 4 = 8, also matching the AOT use_vec8 path).

From my understanding, 32-byte vector loads work correctly on all architectures. On pre-Blackwell GPUs the compiler splits them into two 16-byte load instructions. If #19794 gets merged first, I can update this to use device::kMaxVecBytes introduced by it directly.

@xingsy97 Yes. I think using device::kMaxVecBytes is better.

[xingsy97]

Regression benchmark: 32-byte -> 16-byte vector loads on Blackwell

GPU: NVIDIA GeForce RTX 5060 Laptop GPU (Blackwell, SM120)
Driver: 580.126.20 | CUDA: 12.8 | PyTorch: 2.9.1+cu128
Script: python/sglang/jit_kernel/benchmark/bench_per_token_quant_fp8.py

- m: number of tokens, k: hidden dimension
- 32-byte / 16-byte: max vector load width per memory transaction

per-token-quant-fp8 32 vs 16 byte loads (unit: us):
   m × k            32-byte  16-byte  32->16 regression
   1 × 4096         1.20     1.46     +21.9%
   1 × 7168         1.53     1.67     +9.3%
   8 × 4096         1.39     1.49     +7.2%
   8 × 7168         1.60     1.69     +5.6%
   32 × 4096        1.66     1.71     +3.1%
   32 × 7168        2.09     2.19     +4.7%
   128 × 4096       2.74     3.23     +17.5%
   128 × 7168       4.46     5.04     +13.0%
   512 × 4096       10.22    11.79    +15.3%
   512 × 7168       16.03    18.51    +15.5%
   1024 × 4096      19.07    22.04    +15.6%
   1024 × 7168      31.93    37.79    +18.3%

Conclusion: 32-byte -> 16-byte loads causes 3~22% regression on Blackwell, confirming that 32-byte loads should be kept. Once #19794 merged, this kernel will use device::kMaxVecBytes for automatic arch-based dispatch.

[xingsy97]

Hi @DarkSharpness , I made some improvements to this PR (rebased to latest main branch to include #19794)

Changes

1. Replace hard-coded 32 with device::kMaxVecBytes / sizeof(DType) as discussed earlier.

2. Adopt JIT codebase conventions:

   • packed_t<DType> to process elements in pairs to reduce inner loop count.
   • device::cast<fp32x2_t>(...) instead of static_cast<float> for proper intrinsic dispatch

3. Use warp-local kernel for large batches

• the AOT version dispatches between two kernels depending on token count. The previous code only had the small-batch one (1 token per block, block-level reduce).

• Now added the warp-local variant (8 tokens per block, 1 warp per token, warp-level reduce + optional smem caching), matching AOT behavior.

• Dispatch rule:

  • num_tokens >= sm_count * 2 * 8 -> warp-local,
  • otherwise, small-batch. Smem caching enabled when hidden_dim < 2048 and total smem < 48KB.

4. Updated benchmark and test code to use their new utils methods.

Benchmark & Test

All unit tests pass.

GPU: NVIDIA GeForce RTX 5060 Laptop GPU (Blackwell, SM120, 26 SMs)
CUDA: 12.8 | PyTorch: 2.9.1+cu128
Warp kernel threshold on this GPU: num_tokens >= 416

per-token-quant-fp8 performance (us):

| M x K         | JIT    | AOT    | PyTorch  | JIT to AOT  | kernel      |
|---------------|--------|--------|----------|----------|-------------|
| 1 x 2048      | 1.11   | 1.23   | 14.57    | -9.7%    | small-batch |
| 1 x 4096      | 1.35   | 1.45   | 16.64    | -6.9%    | small-batch |
| 1 x 7168      | 1.66   | 1.71   | 17.25    | -2.9%    | small-batch |
| 16 x 2048     | 1.27   | 1.37   | 20.36    | -7.3%    | small-batch |
| 16 x 4096     | 1.49   | 1.51   | 24.89    | -1.3%    | small-batch |
| 16 x 7168     | 1.70   | 1.81   | 32.24    | -6.1%    | small-batch |
| 64 x 2048     | 1.56   | 1.67   | 22.47    | -6.6%    | small-batch |
| 64 x 4096     | 1.94   | 2.17   | 32.22    | -10.6%   | small-batch |
| 64 x 7168     | 2.73   | 2.82   | 47.48    | -3.2%    | small-batch |
| 128 x 4096    | 2.63   | 3.19   | 48.64    | -17.6%   | small-batch |
| 128 x 7168    | 4.57   | 4.97   | 75.98    | -8.1%    | small-batch |
| 256 x 4096    | 5.61   | 5.70   | 76.17    | -1.6%    | small-batch |
| 256 x 7168    | 9.11   | 9.26   | 156.40   | -1.6%    | small-batch |
| 512 x 2048    | 5.13   | 5.03   | 72.48    | +2.0%    | warp-local  |
| 512 x 4096    | 10.49  | 10.95  | 188.16   | -4.2%    | warp-local  |
| 512 x 7168    | 17.56  | 17.73  | 447.89   | -1.0%    | warp-local  |
| 1024 x 2048   | 9.67   | 9.57   | 183.94   | +1.0%    | warp-local  |
| 1024 x 4096   | 18.39  | 19.15  | 577.07   | -4.0%    | warp-local  |
| 1024 x 7168   | 31.36  | 32.13  | 1452.29  | -2.4%    | warp-local  |

No regression. JIT matches or slightly outperforms AOT across all configs.

[hnyls2002]

Thanks @xingsy97! The JIT tree was relocated from python/sglang/jit_kernel/ into sglang.kernels by #31666, so every path in this diff (including the consumer srt/layers/quantization/fp8_kernel.py, now kernels/ops/quantization/fp8_kernel.py) no longer exists, and the load_jit API changed shape. To be clear the port itself is still wanted - sgl_per_token_quant_fp8 is registered AOT-only at python/sglang/kernels/ops/quantization/__init__.py:20-32. Closing as stale against the new layout; a fresh PR adding kernels/jit/csrc/gemm/per_token_quant_fp8.cuh plus a KernelBackend.JIT registration would be very welcome. Please reopen if I've missed something.