[Kernel] Add JIT activation

[weimin023]

Motivation

Add JIT-compiled CUDA kernels for activation function

Modifications

• Migrate the activation kernels from ahead-of-time (AOT) compilation in sgl_kernel to the JIT compilation framework under python/sglang/jit_kernel/
• Port the CUDA source file (activation.cu) into python/sglang/jit_kernel/csrc/elementwise/activation.cuh with minimal modifications
• Add a Python wrapper (python/sglang/jit_kernel/activation.py)
• Add comprehensive correctness tests (python/sglang/jit_kernel/tests/test_activation.py) to verify JIT kernels with Torch results

Accuracy Tests

pytest /sgl-workspace/sglang/python/sglang/jit_kernel/tests/test_activation.py

platform linux -- Python 3.12.3, pytest-9.0.2, pluggy-1.6.0
rootdir: /sgl-workspace/sglang/python
configfile: pyproject.toml
plugins: anyio-4.12.1, typeguard-4.4.4
collected 48 items

python/sglang/jit_kernel/tests/test_activation.py ................................................ [100%]

=================================================================================================================== 48 passed in 29.37s ====================================================================================================================

Benchmarking and Profiling

Test the accuracy:
python3 -m sglang.test.few_shot_gsm8k --num-questions 200

Accuracy: 0.820
Invalid: 0.000
Latency: 10.294 s
Output throughput: 2829.049 token/s

Benchmark the speed:

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!

[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!

[weimin023]

Hi @BBuf, could you please take a look and let me know which benchmark I should run?
Thanks!

[DarkSharpness]

JIT kernel has not support rocm yet. Maybe just keep the original HIP code first?

[weimin023]

Hi @DarkSharpness, I've kept the HIP code and added a _IS_ROCM guard to allow non-NVIDIA GPUs to use the original AOT kernel. Please let me know if this meets the requirements.

[DarkSharpness]

@weimin023 Thanks for contribution! I adapted from this PR in #21766 (which should be a super-set of this PR) and added you as co-author. Feel free to reopen the PR if something is still missing