feat(nvidia): add ntops RMSNorm backend (#616)

Author: voltjia

CMakeLists.txt

@@ -25,6 +25,8 @@ option(WITH_ASCEND "Enable Ascend backend" OFF)
 
 option(WITH_TORCH "Enable PyTorch C++ backend" OFF)
 
+option(WITH_NINETOOTHED "Enable NineToothed-generated kernels" OFF)
+
 # Default OFF until CANN's `extract_host_stub.py` path handling is fixed for
 # `scikit-build-core` temp-dir builds (triggers `KeyError` on the preprocessed
 # object path). Enable explicitly with `-DBUILD_CUSTOM_KERNEL=ON` when the
@@ -293,6 +295,14 @@ if(_gpu_backend_count GREATER 1)
     message(FATAL_ERROR "`WITH_NVIDIA`, `WITH_ILUVATAR`, `WITH_HYGON`, `WITH_METAX`, `WITH_MOORE`, and `WITH_ASCEND` are mutually exclusive. Build one GPU backend at a time.")
 endif()
 
+if(WITH_NINETOOTHED AND NOT WITH_NVIDIA)
+    message(FATAL_ERROR "`WITH_NINETOOTHED` currently requires `WITH_NVIDIA=ON` because NineToothed AOT uses `caller=\"cuda\"`.")
+endif()
+
+if(WITH_NINETOOTHED)
+    set(NINETOOTHED_PYTHON_EXECUTABLE "" CACHE FILEPATH "Python executable used to run NineToothed code generation")
+endif()
+
 if(WITH_NVIDIA)
     add_compile_definitions(WITH_NVIDIA=1)
     enable_language(CUDA)

scripts/generate_ninetoothed_ops.py

@@ -0,0 +1,87 @@
+import argparse
+import importlib.util
+import pathlib
+import shutil
+import sys
+
+_PROJECT_DIR = pathlib.Path(__file__).resolve().parents[1]
+_OPS_DIR = _PROJECT_DIR / "src" / "ninetoothed" / "ops"
+
+
+def _find_op_modules():
+    return {
+        path.parent.name: path
+        for path in sorted(_OPS_DIR.glob("*/build.py"))
+        if path.is_file()
+    }
+
+
+def _build_manifest(output_dir):
+    return sorted(
+        str(path)
+        for path in pathlib.Path(output_dir).rglob("*.cpp")
+        if not path.name.endswith(".tmp.cpp")
+    )
+
+
+def _write_cmake_manifest(output_dir, sources):
+    manifest_path = pathlib.Path(output_dir) / "manifest.cmake"
+    lines = ["set(INFINIOPS_NINETOOTHED_SOURCES"]
+    lines.extend(f'    "{source}"' for source in sources)
+    lines.append(")")
+    lines.append("")
+    lines.append(f'set(INFINIOPS_NINETOOTHED_INCLUDE_DIRS "{output_dir}")')
+    lines.append("")
+    manifest_path.write_text("\n".join(lines) + "\n")
+
+
+def _load_op_module(op):
+    path = _find_op_modules()[op]
+    sys.path.insert(0, str(path.parent))
+    spec = importlib.util.spec_from_file_location(path.stem, path)
+    module = importlib.util.module_from_spec(spec)
+    assert spec.loader is not None
+    sys.modules[spec.name] = module
+    spec.loader.exec_module(module)
+
+    return module
+
+
+def generate(ops, *, output_dir):
+    op_modules = _find_op_modules()
+    unknown_ops = tuple(op for op in ops if op not in op_modules)
+
+    if unknown_ops:
+        raise ValueError(f"unsupported NineToothed ops: {', '.join(unknown_ops)}")
+
+    output_dir = pathlib.Path(output_dir)
+    shutil.rmtree(output_dir, ignore_errors=True)
+    output_dir.mkdir(parents=True, exist_ok=True)
+
+    for op in ops:
+        module = _load_op_module(op)
+        module.build(output_dir)
+
+    sources = _build_manifest(output_dir)
+    _write_cmake_manifest(output_dir, sources)
+
+    return sources
+
+
+def _parse_args():
+    parser = argparse.ArgumentParser(
+        description="Generate NineToothed operator sources for InfiniOps."
+    )
+    parser.add_argument("--output-dir", required=True)
+    parser.add_argument("--ops", nargs="+", default=tuple(_find_op_modules()))
+
+    return parser.parse_args()
+
+
+def main():
+    args = _parse_args()
+    generate(args.ops, output_dir=args.output_dir)
+
+
+if __name__ == "__main__":
+    main()

scripts/generate_wrappers.py

@@ -1017,11 +1017,13 @@ def _index_impl_headers(impl_roots, scan_dirs):
     return by_operator
 
 
-def _get_all_ops(devices, with_torch=False):
+def _get_all_ops(devices, with_torch=False, with_ninetoothed=False):
     scan_dirs = set(devices)
 
     if with_torch:
         scan_dirs.add("torch")
+    if with_ninetoothed:
+        scan_dirs.add("ninetoothed")
 
     ops = {}
 
@@ -1140,6 +1142,11 @@ def _dispatch_gen_batch_size():
         action="store_true",
         help="Include PyTorch C++ backend implementations.",
     )
+    parser.add_argument(
+        "--with-ninetoothed",
+        action="store_true",
+        help="Include NineToothed backend implementations.",
+    )
 
     args = parser.parse_args()
 
@@ -1159,7 +1166,11 @@ def _dispatch_gen_batch_size():
     if ops_json.exists():
         ops = json.loads(ops_json.read_text())
     else:
-        ops = _get_all_ops(args.devices, with_torch=args.with_torch)
+        ops = _get_all_ops(
+            args.devices,
+            with_torch=args.with_torch,
+            with_ninetoothed=args.with_ninetoothed,
+        )
 
     bind_func_names = []
 

src/CMakeLists.txt

@@ -49,6 +49,39 @@ if(WITH_NVIDIA)
     )
 endif()
 
+if(WITH_NINETOOTHED)
+    find_package(Python COMPONENTS Interpreter REQUIRED)
+
+    if(NINETOOTHED_PYTHON_EXECUTABLE)
+        set(_ninetoothed_python "${NINETOOTHED_PYTHON_EXECUTABLE}")
+    elseif(_TORCH_PYTHON)
+        set(_ninetoothed_python "${_TORCH_PYTHON}")
+    else()
+        set(_ninetoothed_python "${Python_EXECUTABLE}")
+    endif()
+    message(STATUS "NineToothed codegen Python: ${_ninetoothed_python}")
+
+    set(_ninetoothed_output_dir "${CMAKE_CURRENT_BINARY_DIR}/ninetoothed")
+    set(_ninetoothed_generator_args
+        "${PROJECT_SOURCE_DIR}/scripts/generate_ninetoothed_ops.py"
+        --output-dir "${_ninetoothed_output_dir}")
+
+    execute_process(
+        COMMAND "${_ninetoothed_python}" ${_ninetoothed_generator_args}
+        WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}"
+        RESULT_VARIABLE _ninetoothed_generation_result
+    )
+
+    if(NOT _ninetoothed_generation_result EQUAL 0)
+        message(FATAL_ERROR "Generating NineToothed operator sources failed with `${_ninetoothed_python}`. Set `NINETOOTHED_PYTHON_EXECUTABLE` to a Python with `ninetoothed`, `ntops`, `triton`, `sympy`, and CUDA dependencies installed.")
+    endif()
+
+    include("${_ninetoothed_output_dir}/manifest.cmake")
+    target_include_directories(infiniops PRIVATE
+        ${INFINIOPS_NINETOOTHED_INCLUDE_DIRS})
+    target_sources(infiniops PRIVATE ${INFINIOPS_NINETOOTHED_SOURCES})
+endif()
+
 if(WITH_ILUVATAR)
     set(ILUVATAR_PATTERNS
         "native/cuda/*.cc"
@@ -496,6 +529,9 @@ if(GENERATE_CPP_OPERATOR_API OR GENERATE_PYTHON_BINDINGS)
     if(WITH_TORCH)
         list(APPEND GENERATOR_ARGS --with-torch)
     endif()
+    if(WITH_NINETOOTHED)
+        list(APPEND GENERATOR_ARGS --with-ninetoothed)
+    endif()
 
     execute_process(
         COMMAND ${Python_EXECUTABLE} ${PROJECT_SOURCE_DIR}/scripts/generate_wrappers.py ${GENERATOR_ARGS}
@@ -730,6 +766,10 @@ if(GENERATE_PYTHON_BINDINGS)
         ${PROJECT_SOURCE_DIR}/include
         ${PROJECT_SOURCE_DIR}/generated/include
     )
+    if(WITH_NINETOOTHED)
+        target_include_directories(ops PRIVATE
+            ${INFINIOPS_NINETOOTHED_INCLUDE_DIRS})
+    endif()
     target_link_libraries(ops PRIVATE infiniops)
 
     # Cambricon generated dispatch is compiled into the Python extension and