From c6e2d25e3d08cb886b28497cbdc9b68567f11d3a Mon Sep 17 00:00:00 2001 From: Jiacheng Huang Date: Mon, 19 Oct 2026 17:53:44 +0800 Subject: [PATCH] feat: support NineToothed RMSNorm on Ascend --- CMakeLists.txt | 5 +- scripts/generate_ninetoothed_ops.py | 21 ++++- src/CMakeLists.txt | 15 +++ src/ninetoothed/ops/rms_norm/build.py | 61 ++++++++----- src/ninetoothed/ops/rms_norm/rms_norm.h | 116 ++++++++++++++---------- tests/test_generate_ninetoothed_ops.py | 35 +++++++ tests/test_rms_norm.py | 5 +- 7 files changed, 179 insertions(+), 79 deletions(-) diff --git a/CMakeLists.txt b/CMakeLists.txt index 0d8c3e5b3..00246e0e7 100644 --- a/CMakeLists.txt +++ b/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) diff --git a/scripts/generate_ninetoothed_ops.py b/scripts/generate_ninetoothed_ops.py index 387514587..b1e1c715c 100644 --- a/scripts/generate_ninetoothed_ops.py +++ b/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__": diff --git a/src/CMakeLists.txt b/src/CMakeLists.txt index 5902f4aaa..9f3cf105f 100644 --- a/src/CMakeLists.txt +++ b/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() diff --git a/src/ninetoothed/ops/rms_norm/build.py b/src/ninetoothed/ops/rms_norm/build.py index 0f289bddd..7e3d3d122 100644 --- a/src/ninetoothed/ops/rms_norm/build.py +++ b/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, diff --git a/src/ninetoothed/ops/rms_norm/rms_norm.h b/src/ninetoothed/ops/rms_norm/rms_norm.h index 9cd884b45..ec0e47614 100644 --- a/src/ninetoothed/ops/rms_norm/rms_norm.h +++ b/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 weight_sizes; + std::vector weight_strides; + double eps_value = static_cast(eps); + std::int64_t num_normalized_elements = + static_cast(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(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(out.ndim()), 1, + dtype_index, dtype_index, dtype_index); + assert(result == 0 && "NineToothed `RmsNorm` launch failed"); +} + +} // namespace ninetoothed_launcher template <> class Operator : 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 weight_sizes; - std::vector weight_strides; - double eps_value = static_cast(eps); - std::int64_t num_normalized_elements = - static_cast(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(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(stream_), input_tensor, weight_tensor, - eps_tensor, out_tensor, num_normalized_elements_tensor, - static_cast(out.ndim()), 1, dtype_index, dtype_index, dtype_index); + ninetoothed_launcher::LaunchRmsNorm( + input, weight, eps, out, static_cast(stream_)); + } +}; - assert(result == 0 && "NineToothed `RmsNorm` launch failed"); +#ifdef WITH_ASCEND +template <> +class Operator : 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(stream_)); } }; +#endif } // namespace infini::ops - #endif diff --git a/tests/test_generate_ninetoothed_ops.py b/tests/test_generate_ninetoothed_ops.py index 20f8b4861..ebb22f98b 100644 --- a/tests/test_generate_ninetoothed_ops.py +++ b/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" diff --git a/tests/test_rms_norm.py b/tests/test_rms_norm.py index 707f9a85f..59258d79c 100644 --- a/tests/test_rms_norm.py +++ b/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)