Skip to content

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

Draft
fuyou4546 wants to merge 3 commits into
InfiniTensor:masterfrom
fuyou4546:feat/triton-backend-jit
Draft

feat(triton): add JIT backend with add operator #800
fuyou4546 wants to merge 3 commits 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
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.

1 participant