DeepJIT: Header-Only C++20 JIT Runtime for Nvidia CUDA and Huawei Ascend

Developers guyan364, kurisu6912 and LyricZhao released DeepJIT, a header-only C++20 JIT runtime that gives C++ and Python extension authors one shared interface for compiling, caching, loading and launching kernels on NVIDIA CUDA GPUs and Huawei Ascend NPUs. DeepJIT requires CUDA headers 12.4+ and NVCC 12.9+ for the CUDA backend, and CANN with bin/bisheng, bin/ld.lld and Ascend adv_api headers for the Ascend backend, with both backends sharing source and include hashing, in-memory and on-disk caches, and lazy initialization. Two features, cache warmup from history and a Python compilation API, are under development and not yet available.

DeepJIT is a lightweight, header-only C++20 JIT runtime for NVIDIA CUDA GPUs and HUAWEI Ascend 昇腾 NPUs . It gives C++/Python extension authors a shared interface for compiling kernel source at runtime, caching the resulting binaries, loading them onto the device, and launching them with backend-specific options. DeepJIT handles the JIT infrastructure so that kernel libraries can focus on their device code. Both backends share runtime configuration, source and include hashing, in-memory and on-disk caches, and lazy initialization. Kernel source and compiler/launch options remain specific to the selected backend. Main authors: @guyan364 https://github.com/guyan364 , @kurisu6912 https://github.com/kurisu6912 , @LyricZhao https://github.com/LyricZhao . - CUDA and Ascend backends: use deep jit::Runtime<deep jit::CUDA or deep jit::Runtime<deep jit::Ascend with the same compile/load/launch workflow. - Kernel caching: reuse loaded kernels in memory and compiled artifacts on disk. Cache keys account for source, tracked includes, compiler versions, effective compiler options, and an application-provided dependency signature. - Distributed filesystems and shared caches: share one cache directory across users, processes, and nodes to reuse compiled kernels. Both backends support local and distributed filesystems with the required POSIX filesystem semantics; see Shared cache shared-cache for configuration. - Lazy initialization: defer device and compiler discovery until the runtime is first used. - PyTorch integration: use the current PyTorch CUDA or torch npu stream by default, and expose the configured runtime through pybind11 with get jit . - Compilation controls and diagnostics: configure runtime defaults and per-kernel overrides, inspect compilation metadata, and dump CUDA PTX/SASS or Ascend assembly. CUDA also supports a Python post-compilation hook. - Cache warmup from history: use historical cache entries to anticipate kernels that future runs may need and warm up their cache in advance, reducing compilation delays during execution. This feature is under development and is not yet available. - Python compilation API: pass kernel source code directly from Python to compile CUDA or Ascend kernels. This feature is under development and is not yet available. | Backend | Device toolchain and runtime | Integration requirements | |---|---|---| | CUDA | NVCC compiles CUDA source to CUBIN; the CUDA Driver API loads and launches kernels. | CUDA headers 12.4+, NVCC 12.9+, and PyTorch with CUDA support. | | 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. See Integration integration for setup, CUDA cuda for GPU usage, and Ascend ascend for NPU usage. CUDA and Ascend use the same disk-cache implementation. It supports local and distributed filesystems that provide atomic directory rename within a filesystem and file/directory fsync . Builds use unique temporary directories, synchronize their contents, and publish complete entries through an atomic rename. Concurrent processes can compile the same entry and reuse the published result. Multiple users, processes, and nodes can point to the same cache directory: export DJ JIT CACHE DIR=/shared/deep jit Configure directory permissions so participating users can read shared artifacts and writers can create and publish entries under the cache root. As with all DeepJIT caches, use a trusted shared directory. Matching compilation inputs and cache tags allow users to reuse each other's compiled kernels. You can also combine a writable personal cache with a shared lookup cache: export DJ JIT CACHE DIR="$HOME/.dj:/shared/deep jit" DeepJIT searches all roots in order and writes cache misses only to the first root. The shared lookup cache can be read-only. To configure a single consumer library, use its prefix instead, for example MYLIB JIT CACHE DIR . | Path | Contents | |---|---| | include/deep jit/runtime/ /deepseek-ai/DeepJIT/blob/main/include/deep jit/runtime | Shared runtime and configuration. | | include/deep jit/backend/cuda/ /deepseek-ai/DeepJIT/blob/main/include/deep jit/backend/cuda | CUDA compiler, device queries, kernel loading, and launch options. | | include/deep jit/backend/ascend/ /deepseek-ai/DeepJIT/blob/main/include/deep jit/backend/ascend | Ascend compiler/linker integration, device queries, kernel loading, and launch options. | | include/deep jit/cache/ /deepseek-ai/DeepJIT/blob/main/include/deep jit/cache | In-memory and on-disk kernel caches. | | include/deep jit/python api.hpp /deepseek-ai/DeepJIT/blob/main/include/deep jit/python api.hpp | pybind11 registration for a consumer library's runtime. | | tests/ /deepseek-ai/DeepJIT/blob/main/tests | CUDA and Ascend integration tests, example extensions, and device kernels. | The root CMakeLists.txt is for debugging and IDE indexing. Integrate the headers into your own extension as described below; the projects under tests/test cuda proj/ /deepseek-ai/DeepJIT/blob/main/tests/test cuda proj and tests/test ascend proj/ /deepseek-ai/DeepJIT/blob/main/tests/test ascend proj provide working integration examples. Add DeepJIT/include to the include path of the host target, then include: target include directories my target PRIVATE third-party/deep jit/include Include exactly one backend entry header. For CUDA: include <deep jit/backend/cuda/backend.hpp For Ascend: include <deep jit/backend/ascend/backend.hpp The selected header exposes its backend type: using JIT = deep jit::Runtime<deep jit::CUDA ; For the Ascend header, use deep jit::Runtime<deep jit::Ascend instead. create lazy jit delays construction of the runtime until its first use. This also delays device and compiler discovery: inline auto jit = deep jit::create lazy jit<deep jit::CUDA deep jit::Config "/absolute/path/to/my library", "MYLIB" ; If configuration is only known during library initialization, start with an empty lazy object and assign its factory later: include <filesystem include <string include <cutlass/version.h include <deep jit/backend/cuda/backend.hpp namespace my library { inline deep jit::LazyInit<deep jit::Runtime<deep jit::CUDA jit nullptr ; inline void init jit const std::string& library root { const auto library root path = std::filesystem::absolute library root ; const auto include dir = library root path / "include"; jit = deep jit::create lazy jit<deep jit::CUDA deep jit::Config library root path, "MYLIB", "cutlass-" + std::to string CUTLASS VERSION , {include dir}, {"my library/"} ; } } // namespace my library The lazy object must receive a factory before jit- ... or Python get jit is called. deep jit::Config has the following constructor: Config std::filesystem::path python library root, std::string env prefix, std::string extra signature = {}, std::vector<std::filesystem::path include dirs = {}, std::vector<std::string include prefixes = {} ; - python library root resolves relative post hook paths. It must be non-empty and absolute. - env prefix selects the library-specific environment-variable prefix. It must be non-empty and cannot be DJ , which is reserved for global defaults. - extra signature represents dependencies that affect generated code but are not tracked by the include parser. Change it when such a dependency changes, for example "cutlass-" + std::to string CUTLASS VERSION . - include dirs are passed to the compiler and searched by the include parser. Every path must be absolute. - include prefixes select which angle-bracket includes are recursively tracked. For example, "my library/" tracks include <my library/kernel.cuh . Configuration is snapshotted when Runtime is constructed. Do not mutate config , backend , parser , disk cache , or hash base afterward. The supported mutable per-runtime settings are default compiler options and default launch options . Compiler, cache, include, and hook paths, as well as free-form compiler flags, are currently passed through a simple shell command and therefore must not contain whitespace or shell metacharacters. For a pybind11 extension, DeepJIT can register the runtime type and get jit directly. The host library does not need to implement its own get jit binding: include <pybind11/pybind11.h include <deep jit/python api.hpp PYBIND11 MODULE TORCH EXTENSION NAME, module { deep jit::register python api module, my library::jit ; } Python can then obtain the same process-local runtime object: jit = my library. C.get jit Calling get jit initializes the lazy runtime if it has not already been initialized. The CUDA backend requires CUDA headers 12.4 or newer and NVCC 12.9 or newer. It uses NVCC to generate a CUBIN. Loading through compile requires exactly one CUDA kernel; compile without load only builds the artifact and does not perform that check. php const auto kernel = jit- compile "scale", R" extern "C" global void scale float output, const float input, int count { const int index = static cast<int blockIdx.x blockDim.x + threadIdx.x ; if index < count output index = input index 2.0f; } " ; jit- launch kernel, { .grid dim = dim3 count + 255 / 256, 1, 1 , .block dim = dim3 256, 1, 1 , }, output, input, count ; The compile tag must contain only letters, digits, and underscores. compile compiles on a cache miss, loads the CUBIN, and returns a process-cached std::shared ptr<deep jit::cuda::Kernel . compile without load only returns the artifact directory. Unset deep jit::cuda::CompilerOptions fields inherit from the runtime defaults: deep jit::cuda::CompilerOptions options { .optimize level = "3", .fast math = true, .check no spills = true, .arch = jit- device.get arch / use arch family = / false , .extra nvcc flags = {"-DMY OPTION=1"}, .post hook = "hooks/hook 1.py", }; const auto kernel = jit- compile "scale", source, options ; nvcc flags replaces the default free-form NVCC flag list; structured options such as optimize level are generated separately. extra nvcc flags appends per-kernel flags such as -D definitions. To extend the runtime defaults, append flags directly to jit- default compiler options.nvcc flags during initialization. CUDA launch options include the stream, dynamic shared-memory size, grid, block, cluster, cooperative-launch, PDL, and non-portable cluster controls. Unset fields inherit from jit- default launch options , following the same override model as compiler options. Do not clear the initialized runtime defaults back to std::nullopt . Grid and block dimensions are required and must be positive; dynamic shared-memory size cannot be negative. An unset effective stream uses the current PyTorch CUDA stream; an explicitly supplied null stream remains the CUDA default stream. Only one-dimensional clusters are supported cluster dim.y == cluster dim.z == 1 . Enabling nonportable cluster size allowed sets a persistent CUDA function attribute; later launches with the option disabled do not reset that attribute. Compile and load kernels before CUDA Graph capture. Launching an already loaded kernel is capture-compatible, including when the stream is inherited from the current PyTorch stream. Library-wide launch defaults can be changed once during initialization: php jit- default launch options.enable pdl = true; Device information is available from jit- device : php const int num sms = jit- device.get num sms ; const int l2 bytes = jit- device.get num l2 cache bytes ; const int smem bytes = jit- device.get num smem bytes ; const int64 t clock rate = jit- device.get clock rate ; const auto major, minor = jit- device.get arch pair ; const std::string family arch = jit- device.get arch ; const std::string concrete arch = jit- device.get arch false ; A runtime snapshots the current CUDA device when its default architecture is initialized. Use a separate runtime per CUDA device, construct and use it while that device is current, and do not move a loaded kernel between device contexts. CompilerOptions::post hook selects one optional Python file under Config::python library root . For example: my library/hooks/hook 1.py Pass the file's relative path through the compiler options: php const auto kernel = jit- compile "scale", source, { .post hook = "hooks/hook 1.py", } ; post hook inherits from jit- default compiler options like every other compiler option and is unset by default. Only std::nullopt disables it; an empty string is still treated as a configured hook. Once a default hook is set, the current override model cannot disable it for one kernel. The caller must provide a trusted relative path. When set, it runs after NVCC produces the CUBIN and before the artifact is published: cd <temporary-artifact-directory python <absolute-python-library-root /hooks/hook 1.py <absolute-cubin-path The script receives the absolute CUBIN path as its only argument and must modify that file in place. The configured relative post hook path and file-content hash are included in the kernel cache digest; the path is also recorded in meta.json . Hook hashes are cached per thread, so restart the process after changing a hook file. A hook must be deterministic for a given input CUBIN and tracked signature. Imported Python modules, auxiliary files, Python/package versions, environment variables, random state, time, and network data are not discovered automatically; represent every such dependency in extra signature if it can affect output. The CUDA cache digest is built, in order, from: 1. Config::extra signature . 2. The hash of the complete nvcc --version output. 3. The effective compiler flags returned by CompilerOptions::get flags . Paths in Config::include dirs are intentionally excluded, while any -I... placed directly in nvcc flags or extra nvcc flags remains part of the digest. 4. The selected post hook path and file-content hash. 5. The parser digest of the source and its tracked include tree. Each component is prefixed by its fixed-width byte length before it is added to the two-state FNV-1a hash, so boundaries remain unambiguous even for binary strings containing zero bytes. The final digest is a 32-character hexadecimal string. This is a fast cache checksum, not a cryptographic hash. The compile tag is not part of the digest. A disk entry is stored as: <cache-root /cache/<tag .<digest / Cache roots must be trusted. A directory carrying a .committed marker is treated as a completed artifact and its CUBIN may be loaded directly. An entry contains kernel.cu , kernel.cubin , meta.json , .committed , and optional kernel.ptx or kernel.sass files. meta.json records the NVCC command arguments used in the temporary build directory, config fields, compiler information, and effective compiler options, including the selected post hook path. Temporary paths in that command may no longer exist after publication. It does not record the surrounding cd , stderr redirection, post-hook command, or dump commands. dump ptx and dump sass control extra artifacts but do not affect the cache digest. They only run when compilation actually occurs; a pre-existing cache hit is not rebuilt to add missing dumps. The cache key assumes the external compiler environment is stable. Variables such as NVCC PREPEND FLAGS , NVCC APPEND FLAGS , CPATH , host-compiler selection, and transitive third-party header changes are not discovered automatically. Avoid such implicit inputs or represent them in extra signature /explicit compiler flags. Excluding tracked include-directory paths is safe only when moving the same header contents does not change compilation behavior such as embedded FILE strings. The root source is always hashed. An included file is recursively tracked only when both of the following are true: - The directive uses a literal angle-bracket include: include <... . - The included filename starts with one of Config::include prefixes . The parser is intentionally a line-oriented scanner, not a C preprocessor. A tracked directive must use the canonical single-line form include <... whitespace around , include , and the filename is allowed . It does not interpret comments between tokens, macro-expanded includes, backslash continuations, or conditional compilation; an include inside if 0 is still scanned. Tracked files are resolved by checking include dirs in order and using the first matching path. A tracked file that cannot be found is an error. The same rules are applied recursively to includes inside tracked files. Quoted includes such as include "kernel.cuh" are rejected. Macro-based includes are not tracked. Angle-bracket includes that do not match any configured prefix are accepted by the compiler but ignored by the parser. For example, with: deep jit::Config std::filesystem::path "/opt/my library" , "MYLIB", "cutlass-40000", {std::filesystem::path "/opt/my library/include" }, {"my library/"} ; include <my library/kernel.cuh is tracked, while include <cutlass/cutlass.h is not. The CUTLASS version is instead represented by extra signature . Tracked include graphs must not contain cycles. Header digests are cached for the lifetime of a runtime, so recreate the runtime after changing tracked header files. For Config "/absolute/path/to/my library", "MYLIB" , every DeepJIT setting is resolved in this order: MYLIB <SUFFIX DJ <SUFFIX built-in default For example: MYLIB JIT CACHE DIR DJ JIT CACHE DIR $HOME/.dj The library prefix therefore allows one consumer to be configured independently, while the reserved DJ prefix provides process-wide defaults. Boolean values accept true / false , yes / no , or any integer, case-insensitively. Only the library-prefixed and DJ forms are read. An unprefixed variable such as JIT CACHE DIR or JIT DEBUG has no effect. Set JIT environment variables before the first runtime construction normally before the first get jit or jit- ... . Cache roots, compiler selection, C++ standard, and default compiler options are snapshotted then. Compiler-command printing and load-time diagnostics are read again when compile/load runs; changing variables after initialization can therefore produce a mixed configuration and is unsupported. | Suffix | Default | Behavior | |---|---|---| | JIT CACHE DIR | $HOME/.dj | Cache root, or a colon-separated list. All roots are searched in order; misses are compiled into the first root. Empty values or empty list elements are rejected. | | JIT DEBUG | 0 | Enables compiler-command and load diagnostics. CUDA also enables PTXAS output, line info, and PTX/SASS dumps. | | JIT NVCC COMPILER | <discovered-toolkit /bin/nvcc | Overrides the NVCC executable. | | JIT CPP STANDARD | 20 | Selects the C++ standard passed to NVCC as -std=c++<value . | | JIT KERNEL DEBUG INFO | 0 | Adds Bisheng kernel debug information. | | JIT LAUNCH TIMEOUT | 10 | Sets the Ascend kernel launch timeout in seconds; 0 disables it. | | JIT PRINT COMPILER COMMAND | 0 | Prints compiler and disassembler commands. | | JIT PTXAS VERBOSE | 0 | Adds verbose PTXAS output and prints it after compilation. | | JIT CHECK NO SPILLS | 0 | Adds --warn-on-spills and rejects register spills. | | JIT CHECK NO LOCAL MEMORY | 0 | Adds --warn-on-local-memory-usage and rejects any local-memory usage. | | JIT PRINT LOAD TIME | 0 | Prints kernel-binary loading time. | | JIT WITH LINEINFO | 0 | Adds CUDA source line information. | | JIT DUMP ASM | 0 | Generates CUDA PTX/SASS or Ascend assembly artifacts on a cache miss. | | JIT DUMP PTX | 0 | Generates a PTX artifact on a cache miss. | | JIT DUMP SASS | 0 | Generates a SASS artifact on a cache miss. | CUDA toolkit and cache discovery also use these standard environment variables: | Variable | Behavior | |---|---| | HOME | Required for the default $HOME/.dj cache path. | | CUDA HOME | First CUDA toolkit-root candidate. | | CUDA PATH | CUDA toolkit-root fallback when CUDA HOME is unset or empty. | | PATH | Used by which nvcc when neither CUDA root variable identifies a toolkit. | If both CUDA root variables are unset or empty and which nvcc fails, DeepJIT tries /usr/local/cuda . A non-empty but invalid CUDA HOME or CUDA PATH is treated as an error rather than skipped. JIT NVCC COMPILER overrides the executable after CUDA-home discovery. A valid CUDA root must still be discoverable through CUDA HOME , CUDA PATH , PATH , or /usr/local/cuda . SASS dumping additionally requires an executable cuobjdump under that discovered toolkit root. The Ascend backend requires ACL and torch npu headers plus a CANN toolkit containing bin/bisheng and bin/ld.lld . It compiles kernel.asc to kernel.rel.o , links kernel.o , parses the unique .ascend.meta. kernel name, and loads it through ACL. inline auto jit = deep jit::create lazy jit<deep jit::Ascend deep jit::Config "/absolute/path/to/my library", "MYLIB", {}, {"/absolute/path/to/my library/include"}, {"my library/"} ; const auto kernel = jit- compile "scale", source ; jit- launch kernel, {.num blocks = num blocks}, output, input, count ; An unset stream uses the current torch npu stream. num blocks is required. num ubuf bytes controls dynamic UB size, and num launch timeout secs defaults to JIT LAUNCH TIMEOUT . When ASCEND LAUNCH BLOCKING is enabled, DeepJIT synchronizes the device after every launch. deep jit::ascend::CompilerOptions supports the optimization level, dav- architecture, debug information, assembly dumping, and replace/append lists for Bisheng and linker flags. The defaults are -O2 , --cce-aicore-only , VF loop unrolling, and ld.lld -m aicorelinux -Ttext 0 --no-mmap-output-file . <toolkit /aarch64-linux/asc/include/adv api is required and added automatically. Assembly dumping writes Bisheng saved intermediates under the artifact's asm/ directory. Like CUDA's PTX/SASS dumps, it does not change the cache digest and only runs on a cache miss. Device queries are available through jit- device , including get npu arch , get num sms / get num aicore cores , vector and cube core counts, UB size, and L2 size. Toolkit discovery checks ASCEND HOME PATH , ASCEND TOOLKIT HOME , /usr/local/Ascend/ascend-toolkit/latest , and /usr/local/Ascend/cann , in that order.