[jit_kernel] Add fused_qknorm_rope JIT kernel

[Johnsonms]

Motivation

Part of tracking issue #17865 — migrate sgl-kernel AOT kernels to the lightweight python/sglang/jit_kernel/ system.

This PR ports sgl-kernel/csrc/moe/fused_qknorm_rope_kernel.cu to a JIT kernel, matching the existing
sgl_kernel.fused_qk_norm_rope call signature so it can serve as a drop-in replacement.

Key Design Decisions

The original file has no CUTLASS or FlashInfer dependencies — only standard CUDA intrinsics (cuda_bf16.h,
__shfl_xor_sync, __sincosf, etc.). All utilities (packed_as_uint, warp_reduce_sum, div_up, compute_freq_yarn) are inlined
directly in the .cuh file without any heavyweight third-party headers.

One correctness fix is included: the original NeoX active_mask uses (1u << rotary_lanes) - 1, which is undefined behavior
in C++ when rotary_lanes == 32 (the common full-rotary case). The JIT kernel replaces this with 0xffffffffu >> (32 -
rotary_lanes).

Modification

python/sglang/jit_kernel/csrc/elementwise/fused_qknorm_rope.cuh

• fusedQKNormRopeKernel<head_dim, interleave> — one warp per (token, head): packed bf16 loads → RMSNorm (warp-reduce) →
  RoPE (interleave or NeoX, YaRN-aware) → store
• fused_qk_norm_rope host launcher: runtime validation via RuntimeCheck, dispatches to 6 instantiations: head_dim ∈ {64,
  128, 256} × interleave ∈ {true, false}

python/sglang/jit_kernel/fused_qknorm_rope.py

• cache_once/load_jit wrapper; single JIT module (runtime dispatch, no per-head-dim caching)
• fused_qk_norm_rope_out registered as a custom op (mutates qkv in-place)
• fused_qk_norm_rope(..., rotary_dim=None) public API; rotary_dim=None defaults to head_dim

python/sglang/jit_kernel/tests/test_fused_qknorm_rope.py

• 34 correctness tests: vs PyTorch reference across head_dim × num_tokens × is_neox, partial rotary, YaRN scaling,
  default rotary_dim, optional AOT cross-validation

python/sglang/jit_kernel/benchmark/bench_fused_qknorm_rope.py

• JIT vs AOT throughput comparison using triton.testing.perf_report, CI-friendly reduced ranges

Accuracy Tests

34 passed in 2.67s
Correctness: 34/34 passed ✓
Correctness: JIT and AOT are bit-identical across all tested configs
(max_err=0).

python -m pytest python/sglang/jit_kernel/tests/test_fused_qknorm_rope.py -v

Benchmarking and Profiling

python python/sglang/jit_kernel/benchmark/bench_fused_qknorm_rope.py

Performance: JIT matches AOT within noise (<0.5%) across all 15 configurations (1–4096 tokens × head_dim 64/128/256). No meaningful difference.

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

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

[gemini-code-assist]

Summary of Changes

Hello @Johnsonms, 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 integrates a highly optimized, self-contained JIT CUDA kernel for performing fused RMSNorm and Rotary Positional Embeddings on query and key tensors. This enhancement aims to improve the performance and efficiency of attention mechanisms within the framework by consolidating multiple operations into a single, specialized kernel, while also reducing reliance on external, heavier dependencies.

Highlights

• New JIT Kernel for Fused QK-Norm RoPE: A new JIT-compiled CUDA kernel, fused_qknorm_rope, has been introduced, porting the functionality previously found in sgl-kernel/csrc/moe/fused_qknorm_rope_kernel.cu.
• Optimized Fused Operation: The kernel efficiently fuses per-head RMSNorm and Rotary Positional Embeddings (RoPE) for Query (Q) and Key (K) tensors in a single warp-level CUDA pass, operating in-place on a packed QKV tensor.
• Reduced External Dependencies: The implementation is self-contained within a .cuh file, eliminating dependencies on external libraries like CUTLASS or FlashInfer.
• Flexible RoPE Support: The kernel supports various configurations including head_dim values of 64, 128, and 256, both interleave (GPT-J) and NeoX RoPE styles, YaRN frequency scaling, and partial rotary embeddings (rotary_dim <= head_dim).
• Bug Fix for NeoX RoPE: A potential undefined behavior in the NeoX active_mask computation when rotary_lanes == 32 has been addressed and fixed.

🧠 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_fused_qknorm_rope.py
  • Added a new benchmark script to evaluate the performance of the JIT fused_qk_norm_rope against its AOT sgl_kernel counterpart.
  • Included configurations for various num_tokens and head_dim values to cover typical LLM scenarios.
  • Implemented a quick correctness diff check between the JIT and AOT implementations.
• python/sglang/jit_kernel/csrc/elementwise/fused_qknorm_rope.cuh
  • Introduced a new CUDA header file containing the fusedQKNormRopeKernel for fused QK-Norm and RoPE.
  • Implemented warp-level sum reduction and YaRN-aware frequency computation.
  • Provided support for both interleave (GPT-J) and NeoX style RoPE, including a fix for NeoX active_mask UB.
  • Included a host-side TVM-FFI entry point for the kernel, with comprehensive runtime checks for input tensors and parameters.
• python/sglang/jit_kernel/fused_qknorm_rope.py
  • Created a new Python module to wrap the fused_qk_norm_rope JIT CUDA kernel.
  • Defined fused_qk_norm_rope_out as a custom operation that directly calls the CUDA kernel.
  • Provided a user-friendly fused_qk_norm_rope function that handles default rotary_dim behavior and matches the sgl_kernel signature.
• python/sglang/jit_kernel/tests/test_fused_qknorm_rope.py
  • Added a new test suite for the fused_qknorm_rope JIT kernel.
  • Included correctness tests against a pure-PyTorch reference for various head_dim, num_tokens, and is_neox configurations.
  • Verified partial rotary functionality and YaRN scaling behavior.
  • Confirmed that rotary_dim=None correctly defaults to head_dim.
  • Implemented cross-validation tests against the AOT sgl_kernel implementation when available.

Activity

• No specific activity (comments, reviews, or progress updates) was provided in the context for this pull request.

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

The pull request introduces a new JIT kernel for fused QK RMSNorm + RoPE operations, along with benchmarks and correctness tests. The implementation looks solid, and the tests cover various configurations, including partial rotary and YaRN scaling. The code is well-structured and follows good practices for CUDA kernel development and Python integration. The use of static_assert for compile-time checks and RuntimeCheck for runtime validations is good for robustness. The benchmark script provides a clear comparison between the JIT and AOT implementations.

[yuan-luo]

I think the performance comparison doesn't make sense. Because the JIT is faster from the second run onward. We may need to compare that round's statistics.

[yuan-luo]

/tag-and-rerun-ci

[DarkSharpness]

@Johnsonms do you know what's the root cause of the regression?

[Johnsonms]

@Johnsonms do you know what's the root cause of the regression?
Fixed two issues, the regression is gone,
Updated in the PR description, Thanks @yuan-luo @DarkSharpness :

  1. Issue: ~2x perf gap                                                        
  Root Cause: AOT compiled with --use_fast_math; JIT wasn't. powf() in compute_freq_yarn is the bottleneck.                                     
  Fix: Added extra_cuda_cflags=["--use_fast_math"] to load_jit
 
  2. Issue: False MISMATCH (head_dim=256, is_neox=False)
  Root Cause: atol=1e-3 in calculate_diff() was tighter than bfloat16 precision (~7.8e-3/ULP). JIT was actually more correct than AOT vs float32 reference.
  Fix: Raised atol=1e-3 → 1e-2 in calculate_diff(). With --use_fast_math both
    now produce bit-identical outputs (max_err=0).

[yuan-luo]

We also need a python wrapper for this kernel, then we can utilize this kernel and replace the original one.

[Johnsonms]

We also need a python wrapper for this kernel, then we can utilize this kernel and replace the original one.

Done, checked and confirmed

export SGLANG_TORCH_PROFILER_DIR=/scratch/johnson/profile/
python -m sglang.launch_server \
    --model-path Qwen/Qwen3-30B-A3B \
    --enable-fused-qk-norm-rope \
    --tp-size 4

[yuan-luo]

Please fix lint.

[yuan-luo]

Overall LGTM.

[yuan-luo]

Let's wait for the CI.