saymrwulf/onnxruntime
1
0
Fork 0
mirror of https://github.com/saymrwulf/onnxruntime.git synced 2026-07-29 20:14:01 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

[ORTMODULE] Support Register Custom Triton Kernel (#19690)

Browse source 
Add support for registering custom Triton kernel function.
This commit is contained in:
Vincent Wang 2024-02-29 23:03:57 +08:00 committed by GitHub
parent c311d1faf5
commit 937cdd651e
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
2 changed files with 12 additions and 1 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

1
orttraining/orttraining/python/training/ort_triton/__init__.py
View file

@ -9,6 +9,7 @@ from functools import wraps
from onnxruntime.capi import _pybind_state as _C
from .kernel import * # noqa: F403
from .triton_op_executor import register_triton_kernel # noqa: F401
from .triton_op_executor import call_triton_by_name, call_triton_by_onnx, get_config

12
orttraining/orttraining/python/training/ort_triton/triton_op_executor.py
View file

@ -23,6 +23,8 @@ from ._utils import gen_unique_name, next_power_of_2
_DEBUG_MODE = "ORTMODULE_TRITON_DEBUG" in os.environ and int(os.getenv("ORTMODULE_TRITON_DEBUG")) == 1
_CUSTOM_KERNELS = dict()
@functools.lru_cache(None)
def _gen_module_internal(sorted_graph: SortedGraph) -> Tuple[str, str, ModuleType]:
@ -113,7 +115,10 @@ def call_triton_by_name(func_name: str, *tensors, **kwargs):
"""
torch_tensors = [_from_dlpack(tensor) if tensor is not None else None for tensor in tensors]
func = getattr(sys.modules[".".join(__name__.split(".")[:-1])], func_name)
func = getattr(sys.modules[".".join(__name__.split(".")[:-1])], func_name, None)
if func is None:
func = _CUSTOM_KERNELS.get(func_name)
assert func is not None, f"Function {func_name} is not found in the registered kernels."
output = func(*torch_tensors, **kwargs)
if output is not None:
if isinstance(output, tuple):
@ -138,3 +143,8 @@ def call_triton_by_onnx(onnx_key: int, onnx_str: bytes, *tensors):
if isinstance(output, tuple):
return tuple([to_dlpack(tensor) for tensor in output])
return to_dlpack(output)
def register_triton_kernel(fn):
_CUSTOM_KERNELS[fn.__name__] = fn
return fn