feat(triton): add JIT backend with add operator (#800)

* feat(triton): add JIT backend

* refactor(triton): modify naming and move JIT to slot 10

* refactor(triton): replace `CUDA` calls with `InfiniRT` Runtime and Driver APIs

* refactor(triton): simplify JIT backend architecture (#925)

* refactor(triton): simplify JIT backend architecture

* refactor: simplify Triton auto-tuning config

* refactor: minimize Triton JIT operator API

* docs: remove Triton version constraint

* refactor: simplify Triton scalar validation

* refactor: separate Triton platform registration

* refactor: combine Triton Add registration

* style: format Triton JIT headers

---------

Co-authored-by: Jiacheng Huang <45955067+voltjia@users.noreply.github.com>

Author: fuyou4546

CMakeLists.txt

@@ -26,6 +26,8 @@ option(WITH_LINKED "Enable linked third-party operator backends" OFF)
 
 option(WITH_NINETOOTHED "Enable NineToothed-generated kernels" OFF)
 
+option(WITH_TRITON "Enable the Triton JIT backend" OFF)
+
 # Custom `AscendC` kernels under `src/native/ascend/custom/`. `ON` by default
 # so CI and routine dev builds always exercise `implementation_index=1/2`
 # for `RmsNorm` / `AddRmsNorm`. Gated by `WITH_ASCEND` in
@@ -344,10 +346,29 @@ if(WITH_NINETOOTHED)
     set(NINETOOTHED_PYTHON_EXECUTABLE "" CACHE FILEPATH "Python executable used to run NineToothed code generation")
 endif()
 
+if(WITH_TRITON)
+    if(NOT WITH_NVIDIA)
+        message(
+            FATAL_ERROR
+            "`WITH_TRITON` requires `WITH_NVIDIA=ON` because NVIDIA is the only implemented Triton JIT backend."
+        )
+    endif()
+    if(NOT GENERATE_PYTHON_BINDINGS)
+        message(
+            FATAL_ERROR
+            "`WITH_TRITON` requires `GENERATE_PYTHON_BINDINGS=ON` because the runtime compiler is shipped in the Python package."
+        )
+    endif()
+endif()
+
 if(WITH_NVIDIA)
     add_compile_definitions(WITH_NVIDIA=1)
     enable_language(CUDA)
-    find_package(CUDAToolkit REQUIRED)
+    if(WITH_TRITON)
+        find_package(CUDAToolkit 12.0 REQUIRED)
+    else()
+        find_package(CUDAToolkit REQUIRED)
+    endif()
 endif()
 
 # Iluvatar: CUDA-compatible device. CoreX clang++ works for `-x ivcore`

docs/build.md

@@ -19,6 +19,7 @@ entry is `python -m pip install` with CMake options passed through
 | `WITH_TORCH` | Enable PyTorch C++ ATen-backed operators. | `OFF` |
 | `WITH_LINKED` | Enable operators linked from installed third-party libraries. | `OFF` |
 | `WITH_NINETOOTHED` | Enable NineToothed-generated kernels. | `OFF` |
+| `WITH_TRITON` | Enable the NVIDIA Triton JIT backend. Requires Python bindings and CUDA Toolkit 12.0 or newer. | `OFF` |
 | `AUTO_DETECT_DEVICES` | Auto-detect available device files. | `OFF` |
 | `AUTO_DETECT_BACKENDS` | Auto-detect available backend packages. | `OFF` |
 | `GENERATE_OPERATOR_CALL_INSTANTIATIONS` | Generate explicit C++ operator call instantiations. | `ON` |
@@ -51,6 +52,28 @@ python -m pip install .[dev] \
   --config-settings=cmake.define.WITH_NVIDIA=ON
 ```
 
+Install the Triton JIT runtime dependencies, then enable the implementation:
+
+```bash
+python -m pip install torch triton
+python -m pip install . \
+  --config-settings=cmake.define.INFINI_RT_ROOT=/path/to/infini-rt-prefix \
+  --config-settings=cmake.define.WITH_NVIDIA=ON \
+  --config-settings=cmake.define.WITH_TRITON=ON
+```
+
+The InfiniOps JIT bridge and kernel sources are packaged only with the Python
+wheel. It requires CUDA Toolkit 12.0 or newer. Standalone C++ installations do
+not provide this runtime.
+
+Compiled kernels are cached in the platform cache directory. Set
+`INFINI_OPS_TRITON_CACHE_DIR` to override that location.
+
+Python calls with an explicit Triton config construct an operator for that
+call instead of entering the generic operator cache. The compiled kernel and
+auto-tuning result are still cached using the complete Triton config
+identity.
+
 Full builds with both `WITH_NVIDIA=ON` and `WITH_LINKED=ON` include
 `flash_attn_with_kvcache` and require a compatible FlashAttention Python
 distribution in the build environment. The distribution must provide a