Skip to content

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

Draft
fuyou4546 wants to merge 1 commit into
InfiniTensor:masterfrom
fuyou4546:feat/triton-backend-jit
Draft

feat(triton): add JIT backend with add operator #800
fuyou4546 wants to merge 1 commit into
InfiniTensor:masterfrom
fuyou4546:feat/triton-backend-jit

Conversation

@fuyou4546

Copy link
Copy Markdown
Contributor

Summary

  • Triton JIT backend with kernel caching and autotune
  • Add operator JIT implementation (implementation_index=7) via src/triton/ops/add/jit.h
  • scripts/generate_wrappers.py, detects ops using config_t and emits config parameter with inline config-dict parsing
  • src/CMakeLists.txt, compiles JIT infra (jit.cc / compiler.cc) and links pybind11::embed
  • src/config.h, Config::set_extension / extension() for attaching opaque extension data

Motivation

Support JIT compilation for Triton operators. Kernel caching, config passing, and autotune orchestration are implemented in C++. The compilation step bridges to triton.compile via pybind11::embed.

Closes N/A

Type of Change

  • feat — new feature / new operator / new platform
  • fix — bug fix
  • perf — performance improvement (no behavioral change)
  • refactor — code restructuring without behavior change
  • test — adding or fixing tests only
  • docs — documentation only
  • build / ci — build system or CI configuration
  • chore — tooling, formatting, or other non-code changes
  • Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

Platforms Affected

  • CPU (WITH_CPU)
  • NVIDIA (WITH_NVIDIA)
  • Iluvatar (WITH_ILUVATAR)
  • MetaX (WITH_METAX)
  • Cambricon (WITH_CAMBRICON)
  • Moore (WITH_MOORE)
  • Ascend (WITH_ASCEND)
  • PyTorch C++ bindings (WITH_TORCH)
  • Build system / CMake / CI
  • Python bindings / user-facing API

Smoke Test Result

python -m pytest tests -m smoke -q
....................ss..ss....ss....................................ssssss..ss..............                                                  [100%]
78 passed, 14 skipped, 21607 deselected in 2.76s

Test Results on Supported Platforms

Platform Affected Build / Smoke Result Full Result / Notes
NVIDIA Successfully installed InfiniOps-0.1.0 / 78 passed, 14 skipped, 21607 deselected 12237 passed, 9462 skipped, 106 warnings in 356.86s (0:05:56)
Iluvatar
MetaX
Cambricon
Moore
Ascend
Full `pytest` output (optional)
pytest tests -v
···
12237 passed, 9462 skipped, 106 warnings in 356.86s (0:05:56)

Benchmark / Performance Impact

N/A

Notes for Reviewers

  • This PR builds on the AOT backend commits in feat/triton-backend and depends on them being merged first.
  • generate_wrappers.py has additions: _uses_config_extension for detecting ops with config_t, _generate_triton_jit_config_parser for emitting inline config-dict parsing, and conditional uses_config path in _generate_call.
  • Config class (src/config.h) now has a std::shared_ptr<Config> extension_ member with set_extension / extension accessors, used to pass compile configurations from Python bindings to operator implementations.
  • Config dict parsing is emitted inline in the generated binding header.

@fuyou4546 fuyou4546 changed the title Feat/triton backend jit feat(triton): add JIT backend with Add operator Jul 15, 2026
@fuyou4546 fuyou4546 changed the title feat(triton): add JIT backend with Add operator feat(triton): add JIT backend with add operator Jul 15, 2026
@fuyou4546
fuyou4546 force-pushed the feat/triton-backend-jit branch from d2c1db0 to 57b7c00 Compare July 29, 2026 02:57
@fuyou4546
fuyou4546 force-pushed the feat/triton-backend-jit branch from 57b7c00 to f615a7f Compare July 29, 2026 03:05

@voltjia voltjia left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

代码需要遵循 CONTRIBUTING.md,其中 C++ 主要是遵循 Google C++ Style Guide。目前的 naming 之类的好像没有 follow 这些 convention,需要修改一下。

Comment thread src/CMakeLists.txt

if(WITH_TRITON)
target_include_directories(ops PRIVATE
${INFINIOPS_TRITON_INCLUDE_DIRS})

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这里前缀应该是 INFINI_OPS,目前的原则是参数类的都是需要加下划线这种间隔的,只有和链接库名之类相关的才省去。

Comment thread src/CMakeLists.txt
Comment on lines +1086 to +1088
# Ship the JIT compiler and kernel sources so Triton JIT operators
# can compile kernels at runtime. compile.py uses __file__ to
# locate ops/ relative to itself; both must live under triton/.

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

可以的话把注释中的代码用 Markdown 语法括一下,比如 `compile.py``__file__` 之类的,除了此处以外别的地方也检查一下,包括但不限于注释、error message、assert message 等。目前这个规则没有那么严格,如果有注意不到的也没事,不过能注意到的还是改一下吧。

Comment thread src/triton/ops/add/jit.h
#ifndef INFINI_OPS_TRITON_JIT_ADD_H_
#define INFINI_OPS_TRITON_JIT_ADD_H_

#include <cuda.h>

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

triton 下面引入这种平台相关的头文件不太好,因为这样想跨平台就比较难了,要不然就会变成一堆宏。这个地方可以考虑改成 InfiniRT 调用,这样跨平台就交给了 InfiniRT 了,可以保持一致性,而且我看用到的接口目前的 InfiniRT 应该也都提供了。这里也可以像 native 下面的一些算子,把 Backend 作为模板参数,然后通过特化来提供,但是感觉这样对于 triton 下面的东西反而复杂化了。但是总之这里应该尽量不引入具体平台的东西。

Comment thread src/triton/ops/add/jit.h
Comment on lines +1 to +2
#ifndef INFINI_OPS_TRITON_JIT_ADD_H_
#define INFINI_OPS_TRITON_JIT_ADD_H_

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

好像不太对,这个得根据路径来,应该是 INFINI_OPS_TRITON_OPS_ADD_JIT_H

Comment thread src/triton/ops/add/jit.h
namespace infini::ops {

template <>
class Operator<Add, Device::Type::kNvidia, 7> : public Add {

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

这个 slot 先用 10 吧,目前的设定是从 8 开始是特定后端,比如 8 是 PyTorch,9 是 NineToothed,我记得目前应该是 10 还没用到。

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants