Migrate moe_sum_reduce from sgl-kernel AOT to JIT

[Johnsonms]

Motivation

#17865

Modifications

Migrate moe_sum_reduce from the AOT sgl-kernel build to the JIT kernel system (python/sglang/jit_kernel/), following the same pattern established for moe_topk_sigmoid and moe_topk_softmax.

New files:

• python/sglang/jit_kernel/csrc/moe/moe_sum_reduce.cuh — CUDA kernels ported from sgl-kernel/csrc/moe/moe_sum_reduce.cu:
  • moe_sum_reduce_warp_per_token_vec_kernel — vectorized BF16 fast path (16-element uint4 loads)
  • moe_sum_reduce_kernel<T, TOPK> — small-token path with compile-time topk unrolling
  • moe_sum_reduce_kernel_warp_token_topk<T, TOPK, WARPS> — warp-per-token path with compile-time topk unrolling
  • moe_sum_reduce_kernel_general / moe_sum_reduce_kernel_warp_token_general<T, WARPS> — runtime-topk fallbacks
  • tvm-ffi host wrapper moe_sum_reduce replacing PyTorch/ATen tensor handling
• python/sglang/jit_kernel/moe_sum_reduce.py — Python wrapper with the same call signature as sgl_kernel.moe_sum_reduce
• python/sglang/jit_kernel/tests/test_moe_sum_reduce.py — correctness tests
• python/sglang/jit_kernel/benchmark/bench_moe_sum_reduce.py — benchmark

Key implementation notes:

• BF16 vectorized fast path (token_num > 256, hidden_dim % 8 == 0) is gated with if constexpr (std::is_same_v<T,
  __nv_bfloat16>) — zero overhead for fp16/fp32
• Replaced at::opmath_type with explicit float accumulator (same semantics for all three types)
• Replaced AT_DISPATCH_FLOATING_TYPES_AND2 runtime dispatch with JIT-time template instantiation via make_cpp_args(dtype)
• Dropped unused cutlass/array.h include from the original
• Contiguous strides computed from shape; contiguity enforced via RuntimeCheck

Accuracy Tests

677 tests pass, covering:

• All dtypes (fp32, fp16, bf16) across token counts (1–1024), topk values (2, 4, 8, 9, and general fallback with topk=3),
  hidden dims (256, 1024, 4096)
• BF16 vectorized fast path explicitly tested with token_num=512 > 256
• Scale=0 produces all-zeros output
• Cross-validation against sgl_kernel.moe_sum_reduce (AOT): exact match on all tested configs

677 passed in 17.80s

Correctness diff (JIT vs AOT):
tokens= 64 hidden= 4096 topk=4 output=✓ [OK]
tokens= 512 hidden= 7168 topk=8 output=✓ [OK]
tokens= 1024 hidden= 2048 topk=2 output=✓ [OK]

Benchmarking and Profiling

The max deviation across all 24 configs is ±1.8% — this is pure measurement noise, not a real regression. JIT is faster than AOT in 10 configs and slower in 14 configs by the same tiny margin.

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! The target tree python/sglang/jit_kernel/ no longer exists - it was retired into python/sglang/kernels/ (#32072, #33400) - and the tracking issue #17865 is closed. moe_sum_reduce is still AOT-only on main (python/sglang/kernels/aot/python/sgl_kernel/moe.py:85), so a fresh port against python/sglang/kernels/jit/ is still welcome. Closing as obsolete - please reopen if I've missed something.