deepseek-ai/DeepJIT

해결하는 문제

DeepJIT은 GPU 및 NPU 커널을 위한 가벼운 헤더 전용 C++20 런타임을 제공하여 Just-In-Time (JIT) 컴파일 프로세스를 간소화합니다. 커널 라이브러리 개발자가 런타임에 소스 코드를 컴파일하고, 바이너리를 캐시하며, 하드웨어 장치에 로드하기 위한 자체 인프라를 구축할 필요가 없습니다.

작동 방식

DeepJIT은 주로 두 가지 백엔드(NVIDIA CUDA GPU 및 HUAWEI Ascend NPU)에 대해 공통 인터페이스를 제공합니다. 커널의 전체 라이프사이클을 관리합니다: 적절한 툴체인(NVCC for CUDA, Bisheng/ld.lld for Ascend)을 사용하여 소스 코드를 컴파일하고, 결과 바이너리를 디스크 및 메모리에 캐시하여 중복 컴파일을 방지하며, 백엔드별 옵션으로 커널을 실행합니다.

주요 기술적 메커니즘은 다음과 같습니다:

• 해싱 및 캐싱: 소스 코드, 추적된 include 파일, 컴파일러 버전, 구성 옵션을 기반으로 고유한 캐시 키를 생성합니다. 이를 통해 서로 다른 프로세스, 사용자, 심지어 공유 파일 시스템을 통해 다른 노드 간에도 바이너리를 재사용할 수 있습니다.
• include 파서: 특정 각괄호 포함 파일을 추적하는 행 기반 스캐너를 사용하여 종속성 헤더의 변경이 재컴파일을 트리거하도록 합니다.
• 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 스크립트를 실행하여 수정할 수 있습니다.