flagos-ai/libtriton_jit

Purpose and Entry Points

libtriton_jit is a C++20 runtime that invokes Python-defined Triton JIT kernels from C++/libtorch code, reducing Python wrapper, JIT-dispatch, and launch overhead. Its primary API is TritonJITFunction::get_instance(...); operator()(...), declared in include/triton_jit/triton_jit_function.h and implemented in src/triton_jit_function.cpp. Consumer-facing kernel definitions pair Python Triton source with C++ Torch operators, e.g. operators/pointwise/add/add.py and operators/pointwise/add/add_op.cpp.

Runtime Specialization and Kernel Cache

A call to TritonJITFunction captures typed C++ arguments through a variadic template. It combines call-site types, tensor metadata, constexpr values, and backend specialization rules into a full string signature. This determines the compiled-kernel variant and its per-function in-memory cache. The signature/argument utilities and metadata boundary live in include/triton_jit/jit_utils.h, include/triton_jit/kernel_metadata.h, and src/kernel_metadata.cpp.

Python/Triton Compilation Boundary

On a cache miss, the C++ runtime executes scripts/standalone_compile.py, which imports the user’s Python module and delegates compilation to Triton’s Python compiler. The returned compiled artifact path is then loaded by the runtime. Python kernel modules must be directly importable, so relative imports are not viable. This is the key C++/embedded-Python boundary; Triton’s persistent CacheManager owns compiler-wide artifact caching.

Backend-Agnostic Kernel Lifecycle

TritonKernel provides the backend-neutral orchestration: load compiled artifacts, prepare launch state, and launch after compiler-only arguments have been filtered from runtime arguments. Its interface is in include/triton_jit/triton_kernel.h; common policy requirements are defined by include/triton_jit/backend_policy.h. Backend selection is compile-time, rather than a runtime plugin system.

Hardware Backend Policies

Each supported accelerator implements the policy contract for driver/API-specific loading and launching: CUDA, MUSA, NPU/ACL, IX, MACA, MLU, GCU, and HIP-compatible HCU. Representative definitions are include/triton_jit/backends/cuda_backend.h, include/triton_jit/backends/npu_backend.h, and include/triton_jit/backends/hcu_backend.h. NPU-specific ABI argument packing is separated in include/triton_jit/backends/npu_arg_buffer.h. Adding a backend means supplying a conforming policy and CMake integration.

Torch Operator Integration

Operators preprocess Torch tensors, allocate/result-shape outputs, establish device/stream context, invoke Triton kernels, and register callable operations with PyTorch’s dispatcher. Shared operator registration and backend dispatch support are in operators/common/op_registration.h and operators/common/backend_ops.h. Operator families—pointwise, matmul, reduction, indexing, normalization, and fused model kernels—are composed through operators/CMakeLists.txt, with each implementation maintaining a Python kernel plus C++ wrapper boundary.

Build, Distribution, and Dependencies

CMake selects one backend and links Python, libtorch, pybind11, Triton-adjacent dependencies, and vendor SDKs via the root CMakeLists.txt, src/CMakeLists.txt, and backend modules such as cmake/BackendCUDA.cmake. The deliverable is a C++ library with installed headers, compile script, and CMake package metadata. Debian/RPM packaging and CI publish paths are represented by packaging/debian/rules, packaging/rpm/libtriton-jit.spec, and .github/workflows/upload-nexus.yml.