Migrate DeepJIT CUDA integration to the PyTorch stable ABI

[cleonard530]

Summary

• Migrate CUDA stream APIs to the PyTorch 2.10 stable ABI
• Update build configuration, tests, and documentation

The Ascend build is unchanged.

Note: The CUDA extension uses PyTorch’s stable ABI, but it is not CPython ABI-stable because DeepJIT’s GIL helper still depends on pybind11 and CPython-specific APIs. Separate builds are required per Python version.

cc @janeyx99, @Harry-Chen

Testing

• python tests/test_cuda.py : All DeepJIT tests passed

ran torch-abi-audit on minimal extension example (found here 36a98fc) and got

[STABLE  ] [uses-private-api      ] _deep_jit_stable_abi_probe.cpython-312-x86_64-linux-gnu.so  (stable_shim=41, unstable=0)
        cpython violation: PyDict_SetDefault
        cpython violation: PyFrame_GetBack
        cpython violation: PyGILState_Check
        cpython violation: PyInstanceMethod_New
        cpython violation: PyInstanceMethod_Type
        cpython violation: _PyObject_GetDictPtr
        cpython violation: _PyThreadState_UncheckedGet
        cpython violation: _PyType_Lookup 

Most of these stable shims came in from included headers, but some came from #include <deep_jit/backend/cuda/backend.hpp>, and none of them were unstable.

---

Bencmark

These benchmarks compare main against the stable-ABI branch on the SM90 DeepGEMM DeepJIT paths using the same workloads and report both first call and warm call timings. First call measures cold-start end-to-end latency, including JIT compilation, loading, and initial launch overhead in a fresh subprocess with a fresh JIT cache, while warm call measures steady-state per-launch GPU execution after warmup. For evaluating the specific deep_jit launch-path change, warm call is the more direct signal because it isolates runtime launch behavior from one-time compilation costs. The kernels and parameters shapes used for this benchmark are:

• fp8_gemm_nt: a shape (4096, 2048), b shape (7168, 2048)
• grouped_fp8_nt: a shape (3712, 2048), b shape (4, 4096, 2048), grouped_layout shape (3712)
• tf32_hc_prenorm_gemm: a shape (4096, 7168), b shape (24, 7168), num_splits=16
• fp8_mqa_logits: seq=512 kv=4096 h=32 d=128

First Call

Kernel	main (ms)	stable (ms)	diff (ms)
fp8_gemm_nt	3747.508	3915.070	+167.562 (4.47% slower)
grouped_fp8_nt	5453.177	5519.743	+66.566 (1.22% slower)
tf32_hc_prenorm_gemm	1610.871	1584.711	-26.160 (1.62% faster)
fp8_mqa_logits	1542.386	1482.790	-59.596 (-3.86% faster)

Warm Calls (average over 500 calls)

Kernel	main (μs)	stable (μs)	diff (μs)
fp8_gemm_nt	177.9	177.3	-0.6 (0.34% faster)
grouped_fp8_nt	215.5	212.1	-3.4 (1.58% faster)
tf32_hc_prenorm_gemm	62.2	62.3	+0.1 (0.16% slower)
fp8_mqa_logits	49.7	49.8	+0.1 (0.2% slower)

Across the SM90 DeepJIT benchmark cases, the stable-ABI branch runtime is essentially the same as main.

[janeyx99]

Thanks for looking into this! Overall I’m glad we have the shims for the APIs but we should upstream the headeronly bits for better UX.

Could you add a benchmark on a deepJIT workload to validate this change?

[janeyx99]

For more recent torch, we should have a more high level better UX API for this, no?

[cleonard530]

Yes, we definitely have this for torch >= 2.13. I wasn't sure how restrictive we wanted to be here, but seeing your comment below we can update this.

[janeyx99]

Would be good to upstream a high level C++ wrapper API for this one

[janeyx99]

This one too

[janeyx99]

This can be 2.13

[cleonard530]

Could you add a benchmark on a deepJIT workload to validate this change?

Will do!

[cleonard530]

@janeyx99, I added the benchmark results to the description.

[mehmetoguzderin]

Code seems to use different target version at different places, is that intentional? Thanks a lot in advance.

[cleonard530]

Code seems to use different target version at different places, is that intentional? Thanks a lot in advance

No that was not intentional. Thanks for catching that!

[mehmetoguzderin]

Thank you @cleonard530 for the clarification & update!