Migrate DeepJIT CUDA integration to the PyTorch stable ABI by cleonard530 · Pull Request #11 · deepseek-ai/DeepJIT

cleonard530 · 2026-09-11T20:06:42Z

Summary

Migrate CUDA stream APIs to the PyTorch 2.13 stable ABI. This minimum torch version could be brought down to 2.10 with minimal updates if the maintainers would prefer that, and still remain stable. It would just use older APIs.
Update build configuration, tests, and documentation
Added a new flag, DJ_DISABLE_GIL which can be used to make GilScopedRelease a no-op. The reason for this is because if we are using ops registered with TORCH_LIBRARY, the ops invoked through torch.ops releases the GIL already. The current implementation is the only thing in include that is stopping vllm-project/DeepGEMM/15 from using the CPython Stable ABI (abi3) (we remove the <deep_jit/python_api.hpp> directly in that PR, so it didn't effect the CPython Stable ABI).

The Ascend build is unchanged for now, but a followup PR could migrate this to the Torch Stable ABI as well. This might require the torch minimum version staying 2.13 though, as c10_npu::getCurrentNPUStream() would need the shim invoked though stream.nativeHandle(), where auto stream = torch::stable::accelerator::getCurrentStream(device_index);.

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

  -- bundled libs --
    [STABLE  ] [abi3-tagged-no-capi   ] lib.linux-x86_64-cpython-312/_deep_jit_stable_abi_probe.abi3.so  (stable_shim=47, unstable=0)

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.

…uratly

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 · 2026-09-11T21:30:11Z

+inline CUstream get_current_cuda_stream(const int32_t device_index) {
+    void* stream_ptr = nullptr;
+    TORCH_ERROR_CODE_CHECK(
+        aoti_torch_get_current_cuda_stream(device_index, &stream_ptr));


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

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 · 2026-09-11T21:31:23Z

+cudaStream_t get_stream_from_pool(const int32_t device_index) {
+    void* stream = nullptr;
+    TORCH_ERROR_CODE_CHECK(
+        torch_get_cuda_stream_from_pool(false, device_index, &stream));


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

janeyx99 · 2026-09-11T21:31:42Z

+
+    explicit TorchCUDAStreamGuard(cudaStream_t stream, int32_t device_index) {
+        TORCH_ERROR_CODE_CHECK(
+            aoti_torch_create_cuda_stream_guard(


This one too

janeyx99 · 2026-09-11T21:33:15Z

        source_path.write_text(f'#include <{header.as_posix()}>\n', encoding='utf-8')
        command = [
            os.environ.get('CXX', 'c++'), '-std=c++20', '-fsyntax-only', '-Werror',
+            '-DTORCH_TARGET_VERSION=0x020a000000000000', '-DUSE_CUDA',


This can be 2.13

cleonard530 · 2026-09-14T19:10:11Z

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

Will do!

…uda stream instead of C shim directly

cleonard530 · 2026-09-15T13:55:34Z

@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.

Co-authored-by: Mehmet Oguz Derin <mehmetoguzderin@mehmetoguzderin.com>

cleonard530 · 2026-09-15T18:37:45Z

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 · 2026-09-15T18:38:12Z

Thank you @cleonard530 for the clarification & update!

…dRelease used when this macro is set Signed-off-by: Chris Leonard <chleonar@redhat.com>

…lean under -Werror across CUDA, Ascend, and the shared runtime. Signed-off-by: Chris Leonard <chleonar@redhat.com>

janeyx99 · 2026-09-24T21:30:47Z

 | **Ascend** | Bisheng and ld.lld compile and link Ascend kernel source; ACL loads and launches kernels. | CANN with `bin/bisheng`, `bin/ld.lld`, and the Ascend `adv_api` headers; ACL and `torch_npu` headers and runtime. |

-The host environment must provide Linux, a C++20 compiler and standard library with `std::format` support, Python, pybind11, and the dependencies for the selected backend. DeepJIT is intended to be embedded into your extension as a header-only dependency.
+The host environment must provide Linux, a C++20 compiler and standard library with `std::format` support, and the dependencies for the selected backend. The default GIL support also requires Python and pybind11; consumers that manage the GIL themselves can disable it as described under [Integration](#integration). DeepJIT is intended to be embedded into your extension as a header-only dependency. CUDA consumers should compile with `TORCH_TARGET_VERSION=0x020d000000000000` and `USE_CUDA`; this targets PyTorch's stable ABI with PyTorch 2.13 as the minimum runtime version.


@cleonard530 if we use the old cuda stream shim, what's the min runtime version we can target?

We could target 2.10 then (TORCH_TARGET_VERSION=0x020a000000000000)

janeyx99 · 2026-09-25T15:30:25Z

I believe this PR is ready for review @LyricZhao

oraluben

Thanks @cleonard530 !

Overall LGTM. I think we'd prefer the minimal pytorch version to be 2.10 (at least for CUDA), and some minor questions.

oraluben · 2026-09-28T02:05:31Z

 std::shared_ptr<Runtime> process_test_runtime;
 std::shared_ptr<Runtime> gil_test_runtime;

+cudaStream_t get_stream_from_pool(const int32_t device_index) {


If i'm not mistaken, upstream torch is planning to add stable api for these two, so maybe we could move them to a separate util file for future removal?

Yes, we can definitely move this and similar torch functions to a util file!

oraluben · 2026-09-28T02:06:41Z


 pybind11_add_module(_C csrc/python_api.cpp)
+target_compile_definitions(_C PRIVATE
+    TORCH_TARGET_VERSION=0x020d000000000000


It might be better to target 2.10 as minimal requirement as it's feasible. For ascend it's ok to target 2.13 if it's necessary.

Done!

To move to torch 2.10, we needed to replace

auto stream = torch::stable::accelerator::getCurrentStream(device_index); return static_cast<CUstream>(stream.nativeHandle());

with

void* stream_ptr = nullptr; TORCH_ERROR_CODE_CHECK( aoti_torch_get_current_cuda_stream(device_index, &stream_ptr)); return static_cast<CUstream>(stream_ptr);

which requires the macro USE_CUDA to be defined since that method is guarded in PyTorch, so I added this guard in the new include/deep_jit/backend/cuda/stable_torch_utils.h‎ file for it to work

#ifndef USE_CUDA #define DEEP_JIT_UNDEFINE_USE_CUDA #define USE_CUDA #endif #include <torch/csrc/stable/accelerator.h> #include <torch/csrc/stable/c/shim.h> #include <torch/headeronly/util/shim_utils.h> #ifdef DEEP_JIT_UNDEFINE_USE_CUDA #undef USE_CUDA #undef DEEP_JIT_UNDEFINE_USE_CUDA #endif

The final #ifdef block is so the rest of the build is not polluted with the USE_CUDA macro.

oraluben · 2026-09-28T02:13:05Z

+
+namespace deep_jit {
+
+class GilScopedRelease {


I'm not very familiar with the nogil support of pytorch upstream, would you mind elaborate a little bit? Does this mean that DeepJIT itself supports nogil? Or downstream libraries (e.g. deepgemm / deepep) needs extra changes.

When kernels are registered with TORCH_LIBRARY, the operator is registered through PyTorch's dispatcher mechanism (basically becomes just like any other PyTorch op) and can be invoked from Python through torch.ops.<op_name>. When a PyTorch operator is invoked from Python, PyTorch releases the GIL before executing the operator through the dispatcher, so by the time it reaches GilScopedRelease gil_release here, the GIL has already been released and that guard is effectively a no-op. This means that DeepJIT doesn't need to explicitly release the GIL for downstream operators that are registered through TORCH_LIBRARY/STABLE_TORCH_LIBRARY.

Currently, DeepGEMM dispatches through pybind, so the explicit GIL release is still required there. We are working on moving DeepGEMM's registration from pybind to TORCH_LIBRARY (see vllm-project/DeepGEMM#15 and deepseek-ai/DeepGEMM#393). We need that migration to happen before we can remove the explicit GIL release from this path, which is why we added a macro flag DJ_DISABLE_GIL rather than removing the release outright. This allows existing downstream users that still dispatch through pybind to continue working correctly while we migrate.

I'm not as familiar with DeepEP, so if it still relies on pybind or otherwise requires the GIL to be released at this point, that flag would need to remain for it as well, or any other downstream task that aren't registered through TORCH_LIBRARY.

Once the downstream libraries have all moved from pybind to TORCH_LIBRARY/STABLE_TORCH_LIBRARY, and assuming they don't otherwise interact with Python while executing their kernels, we should be able to remove the explicit GilScopedRelease entirely.

Let me know if you have any other questions or need anything clarified. Thanks!

Signed-off-by: Chris Leonard <chleonar@redhat.com>

cleonard530 added 2 commits September 11, 2026 17:42

migrated cuda build to torch stable abi

e90669c

updated kernel.h and main.cpp to match ehat was there before more acc…

a40ccf5

…uratly

janeyx99 reviewed Sep 11, 2026

View reviewed changes

updated TORCH_TARGET_VERSION to 2.13 and used nativeHandle() to get c…

e090381

…uda stream instead of C shim directly

mehmetoguzderin reviewed Sep 15, 2026

View reviewed changes

Comment thread tests/test_cuda.py Resolved

Comment thread CMakeLists.txt Resolved

cleonard530 and others added 2 commits September 15, 2026 14:36

Update tests/test_cuda.py

8662604

Co-authored-by: Mehmet Oguz Derin <mehmetoguzderin@mehmetoguzderin.com>

Update CMakeLists.txt

c2e8a2e

Co-authored-by: Mehmet Oguz Derin <mehmetoguzderin@mehmetoguzderin.com>

cleonard530 added 2 commits September 22, 2026 18:42

added a macro to disable pybind and the gil release. A no-op GilScope…

e4f3e5b

…dRelease used when this macro is set Signed-off-by: Chris Leonard <chleonar@redhat.com>

mark no-op GIL guards as maybe_unused. Keep DJ_DISABLE_GIL=1 builds c…

81ca29e

…lean under -Werror across CUDA, Ascend, and the shared runtime. Signed-off-by: Chris Leonard <chleonar@redhat.com>

This was referenced Sep 23, 2026

Benchmarking DeepGEMM main vs migrate_pybind_to_torch_library_dev branch cleonard530/DeepGEMM#9

Open

Migrate DeepGEMM off of pybind11 and on to TORCH_LIBRARY vllm-project/DeepGEMM#15

Merged

janeyx99 reviewed Sep 24, 2026

View reviewed changes

cleonard530 mentioned this pull request Sep 24, 2026

[RFC]: Enable libtorch-ABI-stable vLLM cuda wheels vllm-project/vllm#26946

Open

1 of 22 tasks

oraluben requested changes Sep 28, 2026

View reviewed changes

Move Torch stable shim helpers into shared header

bad7a43

Signed-off-by: Chris Leonard <chleonar@redhat.com>

oraluben approved these changes Sep 29, 2026

View reviewed changes

LyricZhao approved these changes Sep 30, 2026

View reviewed changes

LyricZhao merged commit 3732a3b into deepseek-ai:main Sep 30, 2026

janeyx99 mentioned this pull request Oct 2, 2026

Add C++ UX wrappers for some CUDA Stream utils stable ABIs pytorch/pytorch#199543

Open


		namespace deep_jit {

		class GilScopedRelease {

Conversation

cleonard530 commented Sep 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Testing

Bencmark

Uh oh!

janeyx99 left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

cleonard530 commented Sep 14, 2026

Uh oh!

cleonard530 commented Sep 15, 2026

Uh oh!

mehmetoguzderin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

cleonard530 commented Sep 15, 2026

Uh oh!

mehmetoguzderin commented Sep 15, 2026

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

janeyx99 commented Sep 25, 2026

Uh oh!

oraluben left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

cleonard530 commented Sep 11, 2026 •

edited

Loading