[jit_kernel] Migrate fused_add_rmsnorm from sgl-kernel AOT to flashinfer JIT

[Johnsonms]

Motivation

#17865
sgl_fused_add_rmsnorm in sgl-kernel wraps flashinfer::norm::FusedAddRMSNorm directly. Since FlashInfer already exposes
fused_add_rmsnorm as a first-class JIT-compiled Python function, maintaining a separate AOT copy adds unnecessary build complexity.
This PR removes the AOT kernel and delegates directly to flashinfer.fused_add_rmsnorm, following the same approach as #18999
(bmm_fp8) and the activation kernel migration.

The existing custom JIT kernel in sglang.jit_kernel.norm is retained for use cases requiring it, while the public
sgl_kernel.fused_add_rmsnorm API delegates to flashinfer which handles all hidden sizes and preserves the enable_pdl parameter.

Modifications

• sgl-kernel/python/sgl_kernel/elementwise.py: Replace fused_add_rmsnorm implementation with a direct call to
  flashinfer.fused_add_rmsnorm, preserving the existing enable_pdl API
• sgl-kernel/csrc/elementwise/fused_add_rms_norm_kernel.cu: Deleted
• sgl-kernel/CMakeLists.txt: Remove fused_add_rms_norm_kernel.cu from build sources
• sgl-kernel/csrc/common_extension.cc: Remove fused_add_rmsnorm op registration
• sgl-kernel/include/sgl_kernel_ops.h: Remove sgl_fused_add_rmsnorm declaration
• sgl-kernel/tests/test_norm.py: Remove test_fused_add_rmsnorm (superseded by JIT tests)
• python/sglang/jit_kernel/tests/test_fused_add_rmsnorm.py: Updated to also compare against AOT kernel (280 tests pass: bf16 × 28
  batch sizes × 10 hidden sizes)
• python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py: Updated to include AOT vs JIT vs FlashInfer comparison

Accuracy Tests

python -m pytest sgl-kernel/tests/test_norm.py -v -s

Benchmarking and Profiling

python python/sglang/jit_kernel/benchmark/bench_fused_add_rmsnorm.py

  Key takeaways:
  - SGL JIT is fastest across most shapes — typically 5–15% faster than
  AOT, 5–20% faster than FlashInfer
  - At hidden=5120, large batch (bs=1024–2048), JIT is dramatically
  faster: ~30–40% over AOT and FlashInfer
  - At hidden=3072, bs=8192, FlashInfer edges ahead — likely a
  memory-bandwidth-bound edge case
  - Some noise at hidden=4096 small batch (GPU scheduling variance), but
  JIT median is still best
  - At very large batch × hidden (memory-bound), all three converge

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]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[hnyls2002]

Thanks @Johnsonms! This has since landed: sgl_kernel.fused_add_rmsnorm now delegates to flashinfer when it is available (python/sglang/kernels/aot/python/sgl_kernel/elementwise.py:155-161), and a native JIT version lives at python/sglang/kernels/jit/csrc/elementwise/fused_add_rmsnorm.cuh. Both paths this PR edits (sgl-kernel/ and python/sglang/jit_kernel/) were also relocated by #31666 / #32072, so the diff no longer applies. Closing as superseded - please reopen if I've missed something.