[cleonard530]

Summary

This has 3 python scripts used for benchmarking the performance between DeepGEMM on main vs migrate_pybind_to_torch_library_dev where I migrated DeepGEMM kernel registration off pybind11 and onto TORCH_LIBRARY. In the runs on main, DeepJIT uses the current version of DeepJIT on main, while the migrate_pybind_to_torch_library_dev branch uses the updated DeepJIT on deepseek-ai/DeepJIT#11, removing the dependency of pybind11 for the GIL release (since torch.ops release the GIL already) and allowing us to build with the CPython Stable ABI (abi3). This branch of DeepJIT also moved the unstable Torch calls to Stable ABI torch calls, which isn't necessary for this first update but is our ultimate goal (see vllm-project/vllm#26946). The only call this changes, is using a stable version of the at::cuda::getCurrentCUDAStream().stream().

The 3 scripts are:

1. scripts/benchmark_kernels.py‎ - script benchmarking all gemm kernels that can run on NVIDIA H200 (plan to run B200 later)
2. scripts/benchmark_sm90_deepjit.py - script to benchmark a select few of the gemm kernels, including the DeepJIT compile times.
3. scripts/benchmark_gil_production.py‎ - Test a single kernel, but testing differnce between pybind11 and TORCH_TENSORS interaction with the GIL.

Results

1. scripts/benchmark_kernels.py‎

The kernel run times where basically identical on main and on migrate_pybind_to_torch_library_dev. You can checkout kernel_benchmark_comparison.txt to see the comparison between average run times for 1000 runs each. The results show ~ +/- 1% difference between runs on each branch.

2. scripts/benchmark_sm90_deepjit.py

The results of this run are in ‎deepjit_main.txt and deepjit_torch_library.txt where each kernel is ran 5 times. Again, the results between main and migrate_pybind_to_torch_library_dev are nearly identical, for cold first call, disk-cache first call, and the warm calls.

3. scripts/benchmark_gil_production.py‎

This script considers the thread management. It basically calls

start worker thread

start timer
call kernel 1,000 times
wait for all GPU work to finish
stop timer

stop worker thread

The script tests no worker, an I/O-like worker, a scheduler-like worker, and a continuously CPU-bound worker. It also measures how much progress the worker makes relative to running alone. Calls through torch.ops release the GIL around operator dispatch, allowing other Python threads to make progress. Under normal I/O-like and scheduler-like activity, this produced little or no foreground latency penalty in these measurements.
The continuously CPU-bound case exposed a tradeoff. Once that thread acquires the GIL, the main thread can be delayed before it can issue its next Python-level operator call. DeepJIT on main also releases the GIL around CUDA dispatch, but only around the kernel call. That shorter release window provides less opportunity for repeated GIL handoffs than releasing it around the complete torch.ops dispatch. Consequently, TORCH_LIBRARY improves the opportunity for competing Python threads to run, but it can increase foreground latency when many short, separate operators compete with a non-yielding CPU-bound Python thread. CUDA graph capture mitigates this case because a captured sequence can be replayed through fewer Python operator calls.

This benchmark intentionally creates an adverse GIL-contention scenario: short kernels, many separate uncaptured calls, and a continuously CPU-bound Python thread in the same process. The behavior is not specific to DeepGEMM; it is a general consideration when replacing a pybind11 binding with a torch.ops binding.

The results can be found ‎gil_production_results_torch_deepjit_main.json and gil_production_results_torch_deepjit_torch_gil.json‎ with a significant difference in median_foreground_us_per_call for the cpu_bound tests.