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

src/ninetoothed/ops/rms_norm/build.py

@@ -0,0 +1,44 @@
+import itertools
+
+import ninetoothed
+import ntops
+
+_BLOCK_SIZES = (256, 512)
+
+_DTYPES = ("float32", "float16", "bfloat16")
+
+_DEFAULT_NDIMS = (2, 3, 4)
+
+_CONFIGS = tuple(
+    (
+        (),
+        {
+            "ndim": ndim,
+            "num_normalized_dims": 1,
+            "input_dtype": dtype,
+            "weight_dtype": dtype,
+            "output_dtype": dtype,
+            "block_size": block_size,
+        },
+        {},
+    )
+    for ndim, dtype, block_size in itertools.product(
+        _DEFAULT_NDIMS,
+        (getattr(ninetoothed, name) for name in _DTYPES),
+        _BLOCK_SIZES,
+    )
+)
+
+
+def build(output_dir):
+    variant_dir = output_dir / "rms_norm"
+    variant_dir.mkdir(parents=True, exist_ok=True)
+    ninetoothed.build(
+        ntops.kernels.rms_norm.premake,
+        _CONFIGS,
+        meta_parameters=("block_size",),
+        caller="cuda",
+        kernel_name="infini_ops_ninetoothed_rms_norm",
+        output_dir=variant_dir,
+        lazy=False,
+    )

src/ninetoothed/ops/rms_norm/rms_norm.h

@@ -0,0 +1,75 @@
+#ifndef INFINI_OPS_NINETOOTHED_RMS_NORM_H_
+#define INFINI_OPS_NINETOOTHED_RMS_NORM_H_
+
+#include <cassert>
+#include <cstdint>
+#include <vector>
+
+#include "base/rms_norm.h"
+#include "data_type.h"
+#include "ninetoothed/tensor.h"
+#include "rms_norm/infini_ops_ninetoothed_rms_norm.h"
+
+namespace infini::ops {
+
+template <>
+class Operator<RmsNorm, Device::Type::kNvidia, 9> : public RmsNorm {
+ public:
+  using RmsNorm::RmsNorm;
+  using RmsNorm::operator();
+
+  void operator()(const Tensor input, const Tensor weight, float eps,
+                  Tensor out) const override {
+    assert(input.dtype() == out.dtype() && out.dtype() == weight.dtype() &&
+           "operator `RmsNorm` requires all input and output tensors to have "
+           "the same dtype");
+    assert(input.shape() == out.shape() &&
+           "NineToothed `RmsNorm` requires input and output tensors with the "
+           "same shape");
+    assert(weight.ndim() == 1 && weight.size(-1) == out.size(-1) &&
+           "NineToothed `RmsNorm` requires a 1D weight matching the last "
+           "dimension");
+    assert(
+        (out.ndim() == 2 || out.ndim() == 3 || out.ndim() == 4) &&
+        "NineToothed `RmsNorm` currently supports rank-2, rank-3, and rank-4 "
+        "tensors");
+
+    std::vector<std::uint64_t> weight_sizes;
+    std::vector<std::int64_t> weight_strides;
+    double eps_value = static_cast<double>(eps);
+    std::int64_t num_normalized_elements =
+        static_cast<std::int64_t>(out.size(-1));
+    std::uint64_t empty_shape[1] = {};
+    std::int64_t empty_strides[1] = {};
+
+    weight_sizes.assign(out.shape().begin(), out.shape().end());
+    weight_strides.assign(out.ndim(), 0);
+    weight_strides.back() =
+        weight.strides().empty() ? 1 : weight.strides().back();
+
+    const int dtype_index = ninetoothed::DataTypeIndex(out.dtype());
+    assert(
+        dtype_index >= 0 &&
+        "NineToothed `RmsNorm` supports only float16, bfloat16, and float32");
+
+    ninetoothed::Tensor input_tensor(input);
+    ninetoothed::Tensor weight_tensor(const_cast<void*>(weight.data()),
+                                      weight_sizes.data(),
+                                      weight_strides.data());
+    ninetoothed::Tensor eps_tensor(eps_value, empty_shape, empty_strides);
+    ninetoothed::Tensor out_tensor(out);
+    ninetoothed::Tensor num_normalized_elements_tensor(
+        num_normalized_elements, empty_shape, empty_strides);
+
+    auto result = launch_infini_ops_ninetoothed_rms_norm(
+        static_cast<NineToothedStream>(stream_), input_tensor, weight_tensor,
+        eps_tensor, out_tensor, num_normalized_elements_tensor,
+        static_cast<int>(out.ndim()), 1, dtype_index, dtype_index, dtype_index);
+
+    assert(result == 0 && "NineToothed `RmsNorm` launch failed");
+  }
+};
+
+}  // namespace infini::ops
+
+#endif

src/ninetoothed/tensor.h

@@ -0,0 +1,64 @@
+#ifndef INFINI_OPS_NINETOOTHED_TENSOR_H_
+#define INFINI_OPS_NINETOOTHED_TENSOR_H_
+
+#include <cstdint>
+#include <type_traits>
+
+#include "data_type.h"
+#include "tensor.h"
+
+namespace infini::ops::ninetoothed {
+
+inline int DataTypeIndex(DataType dtype) {
+  switch (dtype) {
+    case DataType::kFloat16:
+      return 8;
+    case DataType::kBFloat16:
+      return 9;
+    case DataType::kFloat32:
+      return 10;
+    default:
+      return -1;
+  }
+}
+
+class Tensor {
+ public:
+  explicit Tensor(const ::infini::ops::Tensor& tensor)
+      : Tensor(const_cast<void*>(tensor.data()),
+               reinterpret_cast<std::uint64_t*>(
+                   const_cast<::infini::ops::Tensor::Size*>(
+                       tensor.shape().data())),
+               reinterpret_cast<std::int64_t*>(
+                   const_cast<::infini::ops::Tensor::Stride*>(
+                       tensor.strides().data()))) {
+    static_assert(sizeof(::infini::ops::Tensor::Size) == sizeof(std::uint64_t));
+    static_assert(sizeof(::infini::ops::Tensor::Stride) ==
+                  sizeof(std::int64_t));
+    static_assert(std::is_unsigned_v<::infini::ops::Tensor::Size>);
+    static_assert(std::is_signed_v<::infini::ops::Tensor::Stride>);
+  }
+
+  Tensor(void* data, std::uint64_t* shape, std::int64_t* strides)
+      : data_(data), shape_(shape), strides_(strides) {}
+
+  template <typename T>
+  Tensor(T& value, std::uint64_t* shape, std::int64_t* strides)
+      : Tensor(static_cast<void*>(&value), shape, strides) {}
+
+  template <typename NineToothedTensor>
+  operator NineToothedTensor() const {
+    return NineToothedTensor{data_, shape_, strides_};
+  }
+
+ private:
+  void* data_;
+
+  std::uint64_t* shape_;
+
+  std::int64_t* strides_;
+};
+
+}  // namespace infini::ops::ninetoothed
+
+#endif