Summary

• Dispatch hidden size 512 RMSNorm to the existing RMSNormHalfKernel; 1024/1536 stay on the generic CTA path.
• Add a single-warp fast path in rmsnorm_cta_double / rmsnorm_cta_wide that bypasses the CTA shared-memory reduction when kNumWarps == 1.
• Update the RMSNorm dispatch unit test.

H200 Correctness

Environment:

• Host/container: ion8-h200, sglang_bbuf
• Device: NVIDIA H200, CUDA_VISIBLE_DEVICES=4 for unit tests and CUDA_VISIBLE_DEVICES=2/3 for same-shape main/PR benchmarking
• Code under test: local merge-preview with current origin/main (3b3af13d2)

Validation passed:

CUDA_VISIBLE_DEVICES=4 PYTHONPATH=python SGLANG_IS_IN_CI=1 \
TVM_FFI_CACHE_DIR=/tmp/tvm_cache_pr24710_h200 \
pytest -q python/sglang/jit_kernel/tests/test_rmsnorm.py -q

Result: 63 tests passed. Additional BF16 benchmark correctness checks against a PyTorch reference produced matching errors on main and PR: max_abs_error was 0.0, 0.00390625, or 0.0078125 depending on shape.

H200 Benchmark

CUDA event median, current origin/main -> this PR:

The target 512 shape dispatches to RMSNormHalfKernel and preserves correctness. Other rows are included as guard shapes.

Validation

• python -m compileall -q python/sglang/jit_kernel/norm.py python/sglang/jit_kernel/tests/test_rmsnorm.py python/sglang/jit_kernel/csrc/elementwise/rmsnorm.cuh
• Custom H200 RMSNorm BF16 benchmark/correctness script.
• git diff --check origin/main...HEAD

CI States

Latest PR Test (Base): ❌ Run #25919276112
Latest PR Test (Extra): ⚠️ Not enabled -- add run-ci-extra label to opt in.