deepseek-ai/DeepJIT

解决的问题

DeepJIT 提供了一个轻量级的头文件-only C++20 运行时，简化了 GPU 和 NPU 内核的即时（JIT）编译流程。它消除了内核库作者在运行时编译源代码、缓存二进制文件以及将它们加载到硬件设备上时，需要自行构建基础设施的负担。

工作原理

DeepJIT 为两个主要后端（NVIDIA CUDA GPU 和 HUAWEI Ascend NPU）提供统一接口。它管理内核的完整生命周期：使用适当的工具链（CUDA 使用 NVCC，Ascend 使用 Bisheng/ld.lld）编译源代码，将生成的二进制文件缓存到磁盘和内存中以避免重复编译，并使用后端特定选项启动内核。

关键技术机制包括：

• 哈希与缓存：基于源代码、跟踪的头文件、编译器版本和配置选项生成唯一的缓存键。这使得二进制文件可以在不同进程、用户之间，甚至通过共享文件系统在不同节点之间复用。
• 头文件解析器：一个逐行扫描器，用于跟踪特定的尖括号包含文件，以确保依赖头文件的更改会触发重新编译。
• PyTorch 集成：与 PyTorch 的 CUDA 或 torch_npu 流集成，并可通过 pybind11 暴露给 Python。
• 延迟初始化：设备和编译器发现被推迟到运行时实际使用时，以减少启动开销。

适用人群

专为编写高性能设备内核的 C++ 和 Python 扩展作者设计，他们希望为用户提供 JIT 编译体验，而无需实现编译和缓存的底层细节。

特性亮点

• 多后端支持：统一的工作流程，适用于 NVIDIA CUDA 和 HUAWEI Ascend NPU。
• 分布式缓存：支持通过 POSIX 兼容文件系统在多个用户和节点之间共享已编译的内核。
• PyTorch 就绪：与 PyTorch 流无缝集成，并提供 get_jit() Python API。
• 编译诊断：支持导出 PTX/SASS（CUDA）或汇编代码（Ascend）用于调试和检查。
• 编译后钩子：CUDA 后端支持在缓存 CUBIN 文件前运行自定义 Python 脚本进行修改。