[jit_kernel] Add JIT tree_speculative_sampling_target_only kernel

[Johnsonms]

Motivation

Part of the ongoing effort to migrate sgl-kernel AOT kernels to the jit_kernel system (tracking issue #17865). This PR ports sgl-kernel/csrc/speculative/speculative_sampling.cu to a JIT kernel, making tree speculative sampling available without the heavyweight AOT compilation path.

Key Design Decisions

• FlashInfer headers via extra_include_paths: The kernel depends on flashinfer/sampling.cuh for vec_t, SamplingTempStorage, DeviceSamplingFromProb, and dispatch macros (DISPATCH_ALIGNED_VEC_SIZE, DISPATCH_DETERMINISTIC). Following the pattern from rope.py, these headers are included via extra_include_paths at JIT compile time.
• Self-contained kernel: The __global__ kernel and host launcher are copied verbatim from sgl-kernel/csrc/speculative/speculative_sampling.cuh (renamed to TreeSpeculativeSamplingTargetOnlyLauncher to avoid symbol conflicts). A thin tvm-ffi wrapper validates inputs and resolves the CUDA stream.
• Fixed types: The kernel is instantiated with DType=float, IdType=int32_t, IdType2=int64_t matching the AOT implementation.
• API compatibility: The Python wrapper matches sgl_kernel.tree_speculative_sampling_target_only exactly (same keyword argument names and types).

Changes

• python/sglang/jit_kernel/csrc/speculative/speculative_sampling.cuh — CUDA kernel + tvm-ffi host wrapper
• python/sglang/jit_kernel/speculative_sampling.py — Python API with @register_custom_op and @cache_once
• python/sglang/jit_kernel/tests/test_speculative_sampling.py — 13 tests (smoke, AOT cross-validation, boundary)
• python/sglang/jit_kernel/benchmark/bench_speculative_sampling.py — JIT vs AOT throughput benchmark

Test Results

All 13 tests pass:

python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_smoke[1-4-4-32] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_smoke[2-8-5-64] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_smoke[4-16-8-128] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_smoke[1-1-1-32] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_vs_aot[1-4-4-32] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_vs_aot[2-8-5-64] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_vs_aot[4-16-8-128] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_vs_aot[1-1-1-32] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_vs_aot[8-32-8-256] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_accept_all_when_threshold_acc_zero[1] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_accept_all_when_threshold_acc_zero[2] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_accept_all_when_threshold_single_zero[1] PASSED
python/sglang/jit_kernel/tests/test_speculative_sampling.py::test_accept_all_when_threshold_single_zero[2] PASSED
13 passed in 17.09s

Benchmark (JIT vs AOT, µs, A100):
Correctness:

Performance: JIT is at parity or faster than AOT across all 12 configurations.

No regressions. No further tuning needed.

[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 introduces a Just-In-Time (JIT) compiled kernel for tree speculative sampling, migrating an existing Ahead-Of-Time (AOT) kernel to the JIT system. This change aims to streamline the availability of tree speculative sampling by removing the need for a heavy AOT compilation path, thereby improving development and deployment efficiency. The new JIT kernel maintains API compatibility and demonstrates comparable, if not slightly improved, performance against its AOT counterpart.

Highlights

• JIT Kernel Migration: The sgl-kernel AOT tree_speculative_sampling_target_only kernel has been ported to the jit_kernel system, making tree speculative sampling available without heavyweight AOT compilation.
• FlashInfer Integration: FlashInfer headers are now included via extra_include_paths at JIT compile time to provide necessary utilities like vec_t and sampling dispatch macros.
• Self-Contained Implementation: The CUDA kernel and host launcher are copied verbatim from the original sgl-kernel implementation, with a thin tvm-ffi wrapper handling input validation and CUDA stream resolution.
• API Compatibility: The new Python wrapper for the JIT kernel exactly matches the existing sgl_kernel.tree_speculative_sampling_target_only API, ensuring seamless integration.
• Comprehensive Testing and Benchmarking: New tests cover smoke, AOT cross-validation, and boundary conditions, while benchmarks demonstrate the JIT kernel is bitwise identical to AOT and slightly faster (1-2%).

🧠 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_speculative_sampling.py
  • Added a new benchmark script to compare the performance of the JIT tree_speculative_sampling_target_only kernel against its AOT sgl_kernel implementation.
• python/sglang/jit_kernel/csrc/speculative/speculative_sampling.cuh
  • Added the core CUDA kernel and TVM-FFI host wrapper for TreeSpeculativeSamplingTargetOnly, adapted from the sgl-kernel project and integrated with FlashInfer utilities.
• python/sglang/jit_kernel/speculative_sampling.py
  • Added the Python API wrapper for the new JIT tree_speculative_sampling_target_only operation, including register_custom_op and cache_once decorators.
• python/sglang/jit_kernel/tests/test_speculative_sampling.py
  • Added a comprehensive test suite for the JIT tree_speculative_sampling_target_only kernel, covering smoke tests, cross-validation against the AOT version, and boundary conditions for acceptance thresholds.

Activity

• All 13 new tests for the tree_speculative_sampling_target_only kernel passed successfully.
• Benchmarks demonstrated that the new JIT kernel is bitwise identical to the AOT version (when deterministic) and approximately 1-2% faster across various configurations on an A100 GPU.

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 successfully ports the tree_speculative_sampling_target_only kernel from the AOT sgl-kernel to the JIT system. The implementation is clean and follows the established patterns for JIT kernels in the repository. I have identified a few minor issues related to missing header includes in the CUDA source and an inefficient cloning pattern in the benchmark script that could skew performance results.

[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. tree_speculative_sampling_target_only is still AOT-only (python/sglang/kernels/aot/python/sgl_kernel/speculative.py:6), so a fresh port against python/sglang/kernels/jit/ would still be welcome. Closing as obsolete - please reopen if I've missed something.