feat: support NineToothed RMSNorm on Ascend

CMakeLists.txt

@@ -374,12 +374,13 @@ if(_gpu_backend_count GREATER 1)
     message(FATAL_ERROR "`WITH_NVIDIA`, `WITH_ILUVATAR`, `WITH_HYGON`, `WITH_THEAD`, `WITH_METAX`, `WITH_MARS`, `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\"`.")
+if(WITH_NINETOOTHED AND NOT WITH_NVIDIA AND NOT WITH_ASCEND)
+    message(FATAL_ERROR "`WITH_NINETOOTHED` requires either `WITH_NVIDIA=ON` or `WITH_ASCEND=ON`.")
 endif()
 
 if(WITH_NINETOOTHED)
     set(NINETOOTHED_PYTHON_EXECUTABLE "" CACHE FILEPATH "Python executable used to run NineToothed code generation")
+    set(NINETOOTHED_PLATFORM "ascend-910b4" CACHE STRING "NineToothed platform used when WITH_ASCEND is enabled.")
 endif()
 
 if(WITH_TRITON)

scripts/generate_ninetoothed_ops.py

@@ -47,7 +47,7 @@ def _load_op_module(op):
     return module
 
 
-def generate(ops, *, output_dir):
+def generate(ops, *, output_dir, backend="triton", platform=None):
     op_modules = _find_op_modules()
     unknown_ops = tuple(op for op in ops if op not in op_modules)
 
@@ -60,7 +60,7 @@ def generate(ops, *, output_dir):
 
     for op in ops:
         module = _load_op_module(op)
-        module.build(output_dir)
+        module.build(output_dir, backend=backend, platform=platform)
 
     sources = _build_manifest(output_dir)
     _write_cmake_manifest(output_dir, sources)
@@ -74,13 +74,28 @@ def _parse_args():
     )
     parser.add_argument("--output-dir", required=True)
     parser.add_argument("--ops", nargs="+", default=tuple(_find_op_modules()))
+    parser.add_argument(
+        "--backend",
+        choices=("triton", "ascend"),
+        default="triton",
+        help="NineToothed backend used for generated kernels.",
+    )
+    parser.add_argument(
+        "--platform",
+        help="NineToothed platform, such as ascend-910b4.",
+    )
 
     return parser.parse_args()
 
 
 def main():
     args = _parse_args()
-    generate(args.ops, output_dir=args.output_dir)
+    generate(
+        args.ops,
+        output_dir=args.output_dir,
+        backend=args.backend,
+        platform=args.platform,
+    )
 
 
 if __name__ == "__main__":

src/CMakeLists.txt

@@ -190,6 +190,14 @@ if(WITH_NINETOOTHED)
         "${PROJECT_SOURCE_DIR}/scripts/generate_ninetoothed_ops.py"
         --output-dir "${_ninetoothed_output_dir}")
 
+    if(WITH_ASCEND)
+        list(APPEND _ninetoothed_generator_args
+            --backend ascend
+            --platform "${NINETOOTHED_PLATFORM}")
+    else()
+        list(APPEND _ninetoothed_generator_args --backend triton)
+    endif()
+
     execute_process(
         COMMAND "${_ninetoothed_python}" ${_ninetoothed_generator_args}
         WORKING_DIRECTORY "${PROJECT_SOURCE_DIR}"
@@ -203,6 +211,13 @@ if(WITH_NINETOOTHED)
     include("${_ninetoothed_output_dir}/manifest.cmake")
     target_include_directories(infiniops PRIVATE
         ${INFINI_OPS_NINETOOTHED_INCLUDE_DIRS})
+    if(WITH_ASCEND)
+        target_include_directories(infiniops PRIVATE
+            "${ASCEND_HOME}/pkg_inc"
+            "${ASCEND_HOME}/pkg_inc/profiling")
+        target_link_libraries(infiniops PRIVATE
+            "${ASCEND_HOME}/lib64/libruntime.so")
+    endif()
     target_sources(infiniops PRIVATE ${INFINI_OPS_NINETOOTHED_SOURCES})
 endif()
 

src/ninetoothed/ops/rms_norm/build.py

@@ -1,43 +1,54 @@
 import itertools
 
-import ninetoothed
 import ntops
 
-_BLOCK_SIZES = (256, 512)
-
-_DTYPES = ("float32", "float16", "bfloat16")
+import ninetoothed
 
+_BLOCK_SIZES = (256, 512)
 _DEFAULT_NDIMS = (2, 3, 4)
+_DTYPES = {
+    "triton": ("float32", "float16", "bfloat16"),
+    "ascend": ("float32", "float16"),
+}
 
-_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 _configs(backend):
+    return 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[backend]),
+            _BLOCK_SIZES,
+        )
     )
-)
 
 
-def build(output_dir):
+def build(output_dir, *, backend="triton", platform=None):
+    if backend not in _DTYPES:
+        raise ValueError(f"unsupported NineToothed backend: {backend}")
+
+    caller = "npu" if backend == "ascend" else "cuda"
+    platform = platform or ("ascend-910b4" if backend == "ascend" else None)
     variant_dir = output_dir / "rms_norm"
     variant_dir.mkdir(parents=True, exist_ok=True)
     ninetoothed.build(
         ntops.kernels.rms_norm.premake,
-        _CONFIGS,
+        _configs(backend),
         meta_parameters=("block_size",),
-        caller="cuda",
+        backend=backend,
+        platform=platform,
+        caller=caller,
         kernel_name="infini_ops_ninetoothed_rms_norm",
         output_dir=variant_dir,
         lazy=False,

src/ninetoothed/ops/rms_norm/rms_norm.h

@@ -11,65 +11,85 @@
 #include "rms_norm/infini_ops_ninetoothed_rms_norm.h"
 
 namespace infini::ops {
+namespace ninetoothed_launcher {
+
+inline void LaunchRmsNorm(const ::infini::ops::Tensor input,
+                          const ::infini::ops::Tensor weight, float eps,
+                          ::infini::ops::Tensor out, NineToothedStream stream) {
+  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());
+#ifdef WITH_ASCEND
+  assert(out.dtype() != DataType::kBFloat16 &&
+         "Ascend NineToothed RMSNorm does not support bfloat16 yet");
+#endif
+  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);
+  const auto result = launch_infini_ops_ninetoothed_rms_norm(
+      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 ninetoothed_launcher
 
 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);
+    ninetoothed_launcher::LaunchRmsNorm(
+        input, weight, eps, out, static_cast<NineToothedStream>(stream_));
+  }
+};
 
-    assert(result == 0 && "NineToothed `RmsNorm` launch failed");
+#ifdef WITH_ASCEND
+template <>
+class Operator<RmsNorm, Device::Type::kAscend, 9> : public RmsNorm {
+ public:
+  using RmsNorm::RmsNorm;
+  using RmsNorm::operator();
+  void operator()(const Tensor input, const Tensor weight, float eps,
+                  Tensor out) const override {
+    ninetoothed_launcher::LaunchRmsNorm(
+        input, weight, eps, out, static_cast<NineToothedStream>(stream_));
   }
 };
+#endif
 
 }  // namespace infini::ops
-
 #endif

tests/test_generate_ninetoothed_ops.py

@@ -87,6 +87,8 @@ def fake_ntops_premake(*args, **kwargs):
             for block_size in (256, 512)
         )
         assert kwargs["caller"] == "cuda"
+        assert kwargs["backend"] == "triton"
+        assert kwargs["platform"] is None
         assert kwargs["kernel_name"] == "infini_ops_ninetoothed_rms_norm"
         assert kwargs["output_dir"] == tmp_path / "rms_norm"
         assert kwargs["lazy"] is False
@@ -117,3 +119,36 @@ def fake_ntops_premake(*args, **kwargs):
                 },
             )
         ]
+
+
+def test_generate_rms_norm_selects_ascend(monkeypatch):
+    calls = []
+    fake_ninetoothed = types.SimpleNamespace(
+        float32="ninetoothed.float32",
+        float16="ninetoothed.float16",
+        bfloat16="ninetoothed.bfloat16",
+        build=lambda *args, **kwargs: calls.append((args, kwargs)),
+    )
+    module = _load_generator_module(fake_ninetoothed, monkeypatch)
+    fake_ntops = types.SimpleNamespace(
+        kernels=types.SimpleNamespace(
+            rms_norm=types.SimpleNamespace(
+                premake=lambda *args, **kwargs: (args, kwargs, ())
+            )
+        )
+    )
+    monkeypatch.setitem(sys.modules, "ntops", fake_ntops)
+    monkeypatch.setattr(module, "_build_manifest", lambda output_dir: ["kernel.cpp"])
+
+    with tempfile.TemporaryDirectory() as tmpdir:
+        assert module.generate(
+            ["rms_norm"],
+            output_dir=pathlib.Path(tmpdir),
+            backend="ascend",
+            platform="ascend-910b4",
+        ) == ["kernel.cpp"]
+
+    kwargs = calls[0][1]
+    assert kwargs["backend"] == "ascend"
+    assert kwargs["platform"] == "ascend-910b4"
+    assert kwargs["caller"] == "npu"

tests/test_rms_norm.py

@@ -1,7 +1,7 @@
 import infini.ops
 import pytest
-import torch
 
+import torch
 from tests.utils import Payload, empty_strided, get_stream, randn_strided
 
 
@@ -39,6 +39,9 @@ def test_rms_norm(
     rtol,
     atol,
 ):
+    if device == "npu" and implementation_index == 9 and dtype == torch.bfloat16:
+        pytest.skip("Ascend NineToothed RMSNorm does not support bfloat16 yet")
+
     input = randn_strided(input_shape, input_strides, dtype=dtype, device=device)
     weight = randn_strided(weight_shape, weight_strides, dtype=dtype, device=device)
     out = empty_strided(input_shape, out_strides, dtype=dtype, device=device)