refactor(triton): simplify JIT backend architecture

CMakeLists.txt

@@ -25,7 +25,7 @@ option(WITH_TORCH "Enable PyTorch C++ backend" OFF)
 
 option(WITH_NINETOOTHED "Enable NineToothed-generated kernels" OFF)
 
-option(WITH_TRITON "Enable Triton-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`
@@ -336,14 +336,29 @@ if(WITH_NINETOOTHED)
     set(NINETOOTHED_PYTHON_EXECUTABLE "" CACHE FILEPATH "Python executable used to run NineToothed code generation")
 endif()
 
-if(WITH_TRITON AND NOT WITH_NVIDIA)
-    message(FATAL_ERROR "`WITH_TRITON` temporarily requires `WITH_NVIDIA=ON` because the Triton backend temporarily targets CUDA.")
+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

@@ -18,6 +18,7 @@ entry is `python -m pip install` with CMake options passed through
 | `WITH_ASCEND` | Enable the Ascend backend. | `OFF` |
 | `WITH_TORCH` | Enable PyTorch C++ ATen-backed operators. | `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` |
@@ -50,6 +51,27 @@ python -m pip install .[dev] \
   --config-settings=cmake.define.WITH_NVIDIA=ON
 ```
 
+Enable the Triton JIT implementation with the matching optional dependency:
+
+```bash
+python -m pip install .[triton] \
+  --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_TORCH=ON` include
 `flash_attn_with_kvcache` and require a compatible FlashAttention 2.7 Python
 distribution to be installed in the build environment. The distribution must

pyproject.toml

@@ -8,6 +8,7 @@ version = "0.1.0"
 
 [project.optional-dependencies]
 dev = ["pytest", "pytest-cov", "pytest-xdist", "ruff==0.15.22", "torch", "pyyaml"]
+triton = ["torch", "triton>=3.5,<3.6"]
 
 [tool.scikit-build.wheel]
 install-dir = "infini"

scripts/generate_wrappers.py

@@ -1,5 +1,6 @@
 import argparse
 import concurrent.futures
+import dataclasses
 import functools
 import json
 import os
@@ -411,15 +412,34 @@ def _strip_default_argument(param):
     return parts[0].strip()
 
 
+@dataclasses.dataclass(frozen=True)
+class _Implementation:
+    path: pathlib.Path
+    backend: str
+
+
 class _Operator:
-    def __init__(self, name, constructors, calls):
+    def __init__(self, name, constructors, calls, implementations=()):
         self.name = name
 
         self.constructors = constructors
 
         self.calls = calls
 
-        self.impl_paths = []
+        self.implementations = tuple(implementations)
+
+    def has_backend(self, backend):
+        return any(
+            implementation.backend == backend for implementation in self.implementations
+        )
+
+
+def _parameter_type_signature(node):
+    return tuple(
+        " ".join(arg.type.spelling.split())
+        for arg in node.get_arguments()
+        if arg.spelling != "stream"
+    )
 
 
 def _find_optional_tensor_params(op_name):
@@ -582,56 +602,6 @@ def _is_data_type_spelling(spelling):
 
     return spelling.rsplit("::", maxsplit=1)[-1] == "DataType"
 
-def _uses_config_extension(impl_paths):
-    pattern = re.compile(r"\bJitConfig\b")
-    for path in impl_paths:
-        try:
-            if pattern.search(path.read_text()):
-                return True
-        except (OSError, UnicodeDecodeError):
-            pass
-    return False
-
-
-def _generate_triton_jit_config_parser():
-    return textwrap.dedent("""\
-    inline std::shared_ptr<Config> ConfigFromPyDict(const py::dict& config_dict) {
-      if (config_dict.contains("autotune")) {
-        auto config = std::make_shared<AutotuneConfig>();
-        py::dict autotune_dict = config_dict["autotune"].cast<py::dict>();
-        if (autotune_dict.contains("warmup")) config->warmup = autotune_dict["warmup"].cast<unsigned>();
-        if (autotune_dict.contains("rep")) config->rep = autotune_dict["rep"].cast<unsigned>();
-        if (autotune_dict.contains("key")) {
-          for (auto k : autotune_dict["key"].cast<py::list>())
-            config->key.push_back(k.cast<std::string>());
-        }
-        if (autotune_dict.contains("configs")) {
-          for (auto candidate : autotune_dict["configs"].cast<py::list>()) {
-            JitConfig candidate_config;
-            py::dict candidate_dict = candidate.cast<py::dict>();
-            if (candidate_dict.contains("num_warps")) candidate_config.num_warps = candidate_dict["num_warps"].cast<unsigned>();
-            if (candidate_dict.contains("num_stages")) candidate_config.num_stages = candidate_dict["num_stages"].cast<unsigned>();
-            for (auto item : candidate_dict) {
-              std::string key = item.first.cast<std::string>();
-              if (key != "num_warps" && key != "num_stages")
-                candidate_config.constexprs.emplace_back(key, item.second.cast<int>());
-            }
-            config->candidates.push_back(std::move(candidate_config));
-          }
-        }
-        return config;
-      }
-      auto config = std::make_shared<JitConfig>();
-      if (config_dict.contains("num_warps")) config->num_warps = config_dict["num_warps"].cast<unsigned>();
-      if (config_dict.contains("num_stages")) config->num_stages = config_dict["num_stages"].cast<unsigned>();
-      for (auto item : config_dict) {
-        std::string key = item.first.cast<std::string>();
-        if (key != "num_warps" && key != "num_stages")
-          config->constexprs.emplace_back(key, item.second.cast<int>());
-      }
-      return config;
-    }""")
-
 
 def _generate_pybind11(operator):
     optional_tensor_params = _find_optional_tensor_params(operator.name)
@@ -736,7 +706,7 @@ def _generate_arguments(
                 preconverted_tensor_arg is not None
                 and arg.spelling == preconverted_tensor_arg.spelling
             ):
-                args.append(f"std::move({preconverted_tensor_name})")
+                args.append(preconverted_tensor_name)
             elif _is_optional_vector_tensor(arg):
                 args.append(f"VectorOptionalTensorFromPybind11Handle({arg.spelling})")
             elif _is_optional_tensor(arg):
@@ -826,7 +796,7 @@ def _generate_py_args(node):
 
         return ", ".join(parts)
 
-    def _generate_call(op_name, call, method=True, uses_config=False):
+    def _generate_call(op_name, call, method=True, supports_triton_config=False):
         call_params = _generate_params(call)
         call_args = _generate_arguments(call)
 
@@ -848,11 +818,15 @@ def _generate_call(op_name, call, method=True, uses_config=False):
             extra_params = ""
             extra_config_init = ""
             extra_pybind = ""
-            if uses_config:
+            if supports_triton_config:
                 extra_params = ", std::optional<py::dict> config_dict"
                 extra_config_init = (
+                    "    std::unique_ptr<Config> triton_config_ptr;\n"
                     "    if (config_dict.has_value()) {\n"
-                    "      config.set_extension(ConfigFromPyDict(*config_dict));\n"
+                    "      triton_config_ptr = "
+                    "TritonJitConfigFromPyDict(*config_dict);\n"
+                    "      triton_config_ptr->set_implementation_index(\n"
+                    "          config.implementation_index());\n"
                     "    }\n"
                 )
                 extra_pybind = ', py::arg("config") = py::none()'
@@ -869,10 +843,13 @@ def _generate_call(op_name, call, method=True, uses_config=False):
                 call, converted_first_tensor_name
             )
 
-            if uses_config:
+            if supports_triton_config:
                 dispatch = (
-                    f"    auto op = generated_dispatch::Make{symbol_name}(config, {call_args});\n"
-                    f"    (*op)(handle, {call_args});"
+                    "    if (triton_config_ptr) {\n"
+                    f"      auto op = generated_dispatch::Make{symbol_name}(*triton_config_ptr, {call_args});\n"
+                    f"      return (*op)(handle, {call_args});\n"
+                    "    }\n"
+                    f"    return generated_dispatch::Call{symbol_name}(handle, config, {call_args});"
                 )
             else:
                 dispatch = f"    return generated_dispatch::Call{symbol_name}(handle, config, {call_args});"
@@ -943,34 +920,40 @@ def _overload_order_key(node):
     inits = "\n".join(_generate_init(constructor) for constructor in constructors)
     calls = "\n".join(_generate_call(operator.name, call) for call in operator_calls)
 
-    supports_triton = _uses_config_extension(operator.impl_paths)
-    callers = "\n".join(
-        _generate_call(operator.name, call, method=False, uses_config=supports_triton)
+    supports_triton = operator.has_backend("triton")
+    constructor_signatures = {
+        _parameter_type_signature(constructor) for constructor in constructors
+    }
+    configurable_calls = [
+        supports_triton and _parameter_type_signature(call) in constructor_signatures
         for call in operator_calls
-    )
-    if supports_triton:
-        jit_include = (
-            '\n#include "triton/jit/jit.h"\n\n'
-            "namespace infini::ops {\n\n"
-            + _generate_triton_jit_config_parser()
-            + "\n\n}  // namespace infini::ops\n"
+    ]
+    callers = "\n".join(
+        _generate_call(
+            operator.name,
+            call,
+            method=False,
+            supports_triton_config=supports_config,
         )
-    else:
-        jit_include = ""
+        for call, supports_config in zip(operator_calls, configurable_calls)
+    )
+    triton_config_include = (
+        '\n#include "triton/jit/pybind11_config.h"' if any(configurable_calls) else ""
+    )
 
     return f"""#ifndef INFINI_OPS_BINDINGS_{op_name.upper()}_H_
 #define INFINI_OPS_BINDINGS_{op_name.upper()}_H_
 
 #include <pybind11/pybind11.h>
 #include <pybind11/stl.h>
-#include <utility>
+#include <memory>
 
 #include "base/{op_name}.h"
 #include "config.h"
 #include "generated/bindings/generated_dispatch.h"
 #include "handle.h"
 #include "host_range_profiler.h"
-#include "pybind11_utils.h"{jit_include}
+#include "pybind11_utils.h"{triton_config_include}
 
 namespace py = pybind11;
 
@@ -1336,10 +1319,7 @@ def _append_optional_params(prefix, params):
 
     emitted_make_params = set()
 
-    make_nodes = list(operator.constructors)
-    if _uses_config_extension(operator.impl_paths):
-        make_nodes.extend(operator.calls)
-    for node in make_nodes:
+    for node in operator.constructors:
         params = _generate_params(node)
         args = _generate_arguments(node)
         make_params = _append_optional_params("const Config& config", params)
@@ -1751,6 +1731,24 @@ def _matches_scan_dir(impl_path, scan_dirs):
     return any(part in scan_dirs for part in impl_path.parts)
 
 
+def _implementation_backend(path):
+    for backend in ("triton", "torch", "ninetoothed"):
+        if backend in path.parts:
+            return backend
+
+    return "native"
+
+
+def _implementation_from_json(implementation):
+    if isinstance(implementation, str):
+        path = pathlib.Path(implementation)
+        return _Implementation(path, _implementation_backend(path))
+
+    return _Implementation(
+        pathlib.Path(implementation["path"]), implementation["backend"]
+    )
+
+
 _OPERATOR_DECL_RE = re.compile(
     r"\bclass\s+Operator<\s*((?:[A-Za-z_][A-Za-z0-9_]*::)*[A-Za-z_][A-Za-z0-9_]*)\b"
 )
@@ -1774,7 +1772,10 @@ def _index_impl_headers(impl_roots, scan_dirs):
             text = impl_path.read_text()
 
             for match in _OPERATOR_DECL_RE.finditer(text):
-                by_operator.setdefault(match.group(1), []).append(impl_path)
+                implementation = _Implementation(
+                    impl_path, _implementation_backend(impl_path)
+                )
+                by_operator.setdefault(match.group(1), []).append(implementation)
 
     return by_operator
 
@@ -1849,23 +1850,30 @@ def _get_all_ops(devices, with_torch=False, with_ninetoothed=False, with_triton=
             if op_name in ops:
                 continue
 
-            impl_paths = list(
+            implementations = list(
                 impl_headers_by_operator.get(_op_relative_type(op_name), ())
             )
 
-            if not impl_paths:
+            if not implementations:
                 continue
 
-            ops[op_name] = impl_paths
+            ops[op_name] = implementations
 
     return ops
 
 
 def _generate_op_artifacts(item):
-    op_name, impl_paths = item
+    op_name, implementations = item
+    implementations = tuple(implementations)
+    impl_paths = [implementation.path for implementation in implementations]
     extractor = _OperatorExtractor()
-    operator = extractor(op_name)
-    operator.impl_paths = impl_paths
+    parsed_operator = extractor(op_name)
+    operator = _Operator(
+        parsed_operator.name,
+        parsed_operator.constructors,
+        parsed_operator.calls,
+        implementations,
+    )
     header_name = f"{op_name}.h"
     legacy_c_source, legacy_c_header = _generate_legacy_c(operator, impl_paths)
     dispatch_declarations, dispatch_definitions = _generate_generated_dispatch_entries(
@@ -2044,7 +2052,14 @@ def _dispatch_gen_batch_size():
     ops_json = pathlib.Path("ops.json")
 
     if ops_json.exists():
-        ops = json.loads(ops_json.read_text())
+        raw_ops = json.loads(ops_json.read_text())
+        ops = {
+            op_name: [
+                _implementation_from_json(implementation)
+                for implementation in implementations
+            ]
+            for op_name, implementations in raw_ops.items()
+        }
     else:
         ops = _get_all_ops(
             args.devices,