# flagos-ai/libtriton_jit architecture

> Architecture diagram of the GitHub repository flagos-ai/libtriton_jit, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-08-05). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/flagos-ai/libtriton_jit
- Repository: https://github.com/flagos-ai/libtriton_jit

## Overview

#### 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`.

Main entry points: [`include/triton_jit/backend_policy.h`](https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backend_policy.h) (Backend policy contract), [`CMakeLists.txt`](https://github.com/flagos-ai/libtriton_jit/blob/master/CMakeLists.txt) (CMake build), [`src/triton_jit_function.cpp`](https://github.com/flagos-ai/libtriton_jit/blob/master/src/triton_jit_function.cpp) (Function runtime), [`operators/pointwise/add/add_op.cpp`](https://github.com/flagos-ai/libtriton_jit/blob/master/operators/pointwise/add/add_op.cpp) (Add Torch wrapper)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_consumer["Consumer API"]
  node_cpp_consumer["C++ / libtorch consumer<br/>caller"]
end

subgraph group_runtime["JIT Runtime"]
  node_triton_jit_function["TritonJITFunction<br/>public API"]
  node_triton_jit_function_impl["Function runtime<br/>implementation"]
  node_signature_metadata["Signature &amp; metadata<br/>specialization utilities<br/>[jit_utils.h]"]
  node_kernel_metadata["Kernel metadata<br/>metadata boundary<br/>[kernel_metadata.h]"]
  node_kernel_metadata_impl["Metadata implementation"]
  node_triton_kernel["TritonKernel<br/>kernel lifecycle<br/>[triton_kernel.h]"]
end

subgraph group_compiler["Python Compilation"]
  node_standalone_compile["Standalone compiler<br/>Python compile bridge"]
  node_python_kernel["Python Triton kernel<br/>kernel source<br/>[add.py]"]
  node_triton_cache[("Triton CacheManager<br/>persistent artifact cache")]
end

subgraph group_backends["Backend Policies"]
  node_backend_policy["Backend policy contract<br/>policy interface<br/>[backend_policy.h]"]
  node_cuda_backend["CUDA policy<br/>backend policy<br/>[cuda_backend.h]"]
  node_npu_backend["NPU / ACL policy<br/>backend policy<br/>[npu_backend.h]"]
  node_npu_arg_buffer["NPU argument buffer<br/>ABI packing<br/>[npu_arg_buffer.h]"]
  node_hcu_backend["HCU / HIP policy<br/>backend policy<br/>[hcu_backend.h]"]
end

subgraph group_operators["Torch Operators"]
  node_torch_wrapper["Add Torch wrapper<br/>C++ operator<br/>[add_op.cpp]"]
  node_operator_registration["Operator registration<br/>PyTorch dispatcher support<br/>[op_registration.h]"]
  node_operator_families["Operator families<br/>operator composition<br/>[CMakeLists.txt]"]
end

subgraph group_build["Build &amp; Distribution"]
  node_root_build["CMake build<br/>build orchestration<br/>[CMakeLists.txt]"]
  node_cuda_build["CUDA CMake module<br/>backend integration<br/>[BackendCUDA.cmake]"]
  node_package_publish["Package publishing<br/>CI distribution<br/>[upload-nexus.yml]"]
end

node_cpp_consumer -->|"calls"| node_triton_jit_function
node_triton_jit_function -->|"implements"| node_triton_jit_function_impl
node_triton_jit_function_impl -->|"specializes from arguments"| node_signature_metadata
node_signature_metadata -->|"uses metadata"| node_kernel_metadata
node_kernel_metadata -->|"implemented by"| node_kernel_metadata_impl
node_triton_jit_function_impl -->|"cache miss compiles"| node_standalone_compile
node_standalone_compile -->|"imports"| node_python_kernel
node_standalone_compile -->|"reads and writes artifacts"| node_triton_cache
node_triton_jit_function_impl -->|"loads cached variant"| node_triton_kernel
node_triton_kernel -->|"orchestrates through"| node_backend_policy
node_backend_policy -->|"implemented by"| node_cuda_backend
node_backend_policy -->|"implemented by"| node_npu_backend
node_backend_policy -->|"implemented by"| node_hcu_backend
node_npu_backend -->|"packs ABI arguments"| node_npu_arg_buffer
node_torch_wrapper -->|"invokes"| node_triton_jit_function
node_torch_wrapper -.->|"pairs with"| node_python_kernel
node_torch_wrapper -->|"registers"| node_operator_registration
node_operator_families -->|"builds"| node_torch_wrapper
node_root_build -->|"includes"| node_operator_families
node_root_build -->|"selects backend integration"| node_cuda_build
node_root_build -->|"chooses implementation at compile time"| node_backend_policy
node_root_build -->|"produces publishable artifacts"| node_package_publish

click node_triton_jit_function "https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/triton_jit_function.h"
click node_triton_jit_function_impl "https://github.com/flagos-ai/libtriton_jit/blob/master/src/triton_jit_function.cpp"
click node_signature_metadata "https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/jit_utils.h"
click node_kernel_metadata "https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/kernel_metadata.h"
click node_kernel_metadata_impl "https://github.com/flagos-ai/libtriton_jit/blob/master/src/kernel_metadata.cpp"
click node_triton_kernel "https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/triton_kernel.h"
click node_standalone_compile "https://github.com/flagos-ai/libtriton_jit/blob/master/scripts/standalone_compile.py"
click node_python_kernel "https://github.com/flagos-ai/libtriton_jit/blob/master/operators/pointwise/add/add.py"
click node_backend_policy "https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backend_policy.h"
click node_cuda_backend "https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backends/cuda_backend.h"
click node_npu_backend "https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backends/npu_backend.h"
click node_npu_arg_buffer "https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backends/npu_arg_buffer.h"
click node_hcu_backend "https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backends/hcu_backend.h"
click node_torch_wrapper "https://github.com/flagos-ai/libtriton_jit/blob/master/operators/pointwise/add/add_op.cpp"
click node_operator_registration "https://github.com/flagos-ai/libtriton_jit/blob/master/operators/common/op_registration.h"
click node_operator_families "https://github.com/flagos-ai/libtriton_jit/blob/master/operators/CMakeLists.txt"
click node_root_build "https://github.com/flagos-ai/libtriton_jit/blob/master/CMakeLists.txt"
click node_cuda_build "https://github.com/flagos-ai/libtriton_jit/blob/master/cmake/BackendCUDA.cmake"
click node_package_publish "https://github.com/flagos-ai/libtriton_jit/blob/master/.github/workflows/upload-nexus.yml"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_cpp_consumer toneBlue
class node_triton_jit_function,node_triton_jit_function_impl,node_signature_metadata,node_kernel_metadata,node_kernel_metadata_impl,node_triton_kernel toneAmber
class node_standalone_compile,node_python_kernel,node_triton_cache toneMint
class node_backend_policy,node_cuda_backend,node_npu_backend,node_npu_arg_buffer,node_hcu_backend toneRose
class node_torch_wrapper,node_operator_registration,node_operator_families toneIndigo
class node_root_build,node_cuda_build,node_package_publish toneTeal
```

## Components

### Consumer API

- **C++ / libtorch consumer**: caller

### JIT Runtime

- [TritonJITFunction](https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/triton_jit_function.h): public API · Typed call capture, specialization, and per-function variant cache. · `include/triton_jit/triton_jit_function.h`
- [Function runtime](https://github.com/flagos-ai/libtriton_jit/blob/master/src/triton_jit_function.cpp): implementation · `src/triton_jit_function.cpp`
- [Signature & metadata](https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/jit_utils.h): specialization utilities · Builds signatures from types, tensors, constexpr values, and backend rules. · `include/triton_jit/jit_utils.h`
- [Kernel metadata](https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/kernel_metadata.h): metadata boundary · `include/triton_jit/kernel_metadata.h`
- [Metadata implementation](https://github.com/flagos-ai/libtriton_jit/blob/master/src/kernel_metadata.cpp): implementation · `src/kernel_metadata.cpp`
- [TritonKernel](https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/triton_kernel.h): kernel lifecycle · Loads artifacts, filters compiler-only arguments, prepares and launches. · `include/triton_jit/triton_kernel.h`

### Python Compilation

- [Standalone compiler](https://github.com/flagos-ai/libtriton_jit/blob/master/scripts/standalone_compile.py): Python compile bridge · Imports the user module and invokes Triton's Python compiler. · `scripts/standalone_compile.py`
- [Python Triton kernel](https://github.com/flagos-ai/libtriton_jit/blob/master/operators/pointwise/add/add.py): kernel source · Must be directly importable by the compiler process. · `operators/pointwise/add/add.py`
- **Triton CacheManager**: persistent artifact cache · Owns compiler-wide compiled-artifact caching.

### Backend Policies

- [Backend policy contract](https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backend_policy.h): policy interface · `include/triton_jit/backend_policy.h`
- [CUDA policy](https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backends/cuda_backend.h): backend policy · `include/triton_jit/backends/cuda_backend.h`
- [NPU / ACL policy](https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backends/npu_backend.h): backend policy · `include/triton_jit/backends/npu_backend.h`
- [NPU argument buffer](https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backends/npu_arg_buffer.h): ABI packing · `include/triton_jit/backends/npu_arg_buffer.h`
- [HCU / HIP policy](https://github.com/flagos-ai/libtriton_jit/blob/master/include/triton_jit/backends/hcu_backend.h): backend policy · `include/triton_jit/backends/hcu_backend.h`

### Torch Operators

- [Add Torch wrapper](https://github.com/flagos-ai/libtriton_jit/blob/master/operators/pointwise/add/add_op.cpp): C++ operator · Prepares tensors and invokes the specialized kernel. · `operators/pointwise/add/add_op.cpp`
- [Operator registration](https://github.com/flagos-ai/libtriton_jit/blob/master/operators/common/op_registration.h): PyTorch dispatcher support · `operators/common/op_registration.h`
- [Operator families](https://github.com/flagos-ai/libtriton_jit/blob/master/operators/CMakeLists.txt): operator composition · Builds pointwise, matmul, reduction, indexing, normalization, and fused kernels. · `operators/CMakeLists.txt`

### Build & Distribution

- [CMake build](https://github.com/flagos-ai/libtriton_jit/blob/master/CMakeLists.txt): build orchestration · Selects one backend and links Python, libtorch, pybind11, Triton dependencies, and SDKs. · `CMakeLists.txt`
- [CUDA CMake module](https://github.com/flagos-ai/libtriton_jit/blob/master/cmake/BackendCUDA.cmake): backend integration · `cmake/BackendCUDA.cmake`
- [Package publishing](https://github.com/flagos-ai/libtriton_jit/blob/master/.github/workflows/upload-nexus.yml): CI distribution · `.github/workflows/upload-nexus.yml`

## Connections

- C++ / libtorch consumer → TritonJITFunction: calls
- TritonJITFunction → Function runtime: implements
- Function runtime → Signature & metadata: specializes from arguments
- Signature & metadata → Kernel metadata: uses metadata
- Kernel metadata → Metadata implementation: implemented by
- Function runtime → Standalone compiler: cache miss compiles
- Standalone compiler → Python Triton kernel: imports
- Standalone compiler → Triton CacheManager: reads and writes artifacts
- Function runtime → TritonKernel: loads cached variant
- TritonKernel → Backend policy contract: orchestrates through
- Backend policy contract → CUDA policy: implemented by
- Backend policy contract → NPU / ACL policy: implemented by
- Backend policy contract → HCU / HIP policy: implemented by
- NPU / ACL policy → NPU argument buffer: packs ABI arguments
- Add Torch wrapper → TritonJITFunction: invokes
- Add Torch wrapper → Python Triton kernel: pairs with
- Add Torch wrapper → Operator registration: registers
- Operator families → Add Torch wrapper: builds
- CMake build → Operator families: includes
- CMake build → CUDA CMake module: selects backend integration
- CMake build → Backend policy contract: chooses implementation at compile time
- CMake build → Package publishing: produces publishable artifacts