feat(nvidia): add ntops rms norm backend by voltjia · Pull Request #616 · InfiniTensor/InfiniOps

voltjia · 2026-05-20T02:04:35Z

Summary

Add optional NineToothed code generation for NVIDIA behind WITH_NINETOOTHED, driven by ntops.kernels.rms_norm.premake and ninetoothed.build.
Add an NVIDIA slot 9 RmsNorm implementation that launches generated NineToothed kernels and adapts InfiniOps' 1D weight tensor to the expanded tensor view expected by ntops.
Keep scripts/generate_ninetoothed_ops.py as a build entrypoint and move the operator-specific codegen implementation under src/native/ninetoothed.
Add a small generator unit test covering the ntops premake path and generated CMake manifest behavior.

Motivation

ntops already provides the operator premake/application logic, so InfiniOps should only own the integration layer: selecting supported ranks/dtypes, running ninetoothed.build, compiling the generated sources, and dispatching through the existing implementation index mechanism.

This PR starts with RmsNorm because it is currently available in InfiniTensor/ntops; Swiglu is intentionally left out until ntops exposes a corresponding premake.

Closes # N/A

Type of Change

feat - new feature / new operator / new platform
fix - bug fix
perf - performance improvement (no behavioral change)
refactor - code restructuring without behavior change
test - adding or fixing tests only
docs - documentation only
build / ci - build system or CI configuration
chore - tooling, formatting, or other non-code changes
Breaking change (requires a ! in the Conventional Commits prefix or a BREAKING CHANGE: footer)

Platforms Affected

Test Results on Supported Platforms

Platform	Built	`pytest` Result	Notes / Hardware
NVIDIA	Yes	Targeted smoke passed	Remote `ssh nvidia`, `infiniops-ci/nvidia:latest`; built `ops` with `WITH_NVIDIA=ON`, `WITH_NINETOOTHED=ON`, `GENERATE_PYTHON_BINDINGS=ON`; ran `RmsNorm` slot 9 smoke vs. PyTorch on rank-2 and rank-3 `float32`, max error `2.38e-07`, `allclose=True`.
Iluvatar	No	Not run	Not touched; no hardware used for this PR.
MetaX	No	Not run	Not touched; no hardware used for this PR.
Cambricon	No	Not run	Not touched; no hardware used for this PR.
Moore	No	Not run	Not touched; no hardware used for this PR.
Ascend	No	Not run	Not touched; no hardware used for this PR.

Full `pytest` output (optional)

$ python3 -m py_compile scripts/generate_ninetoothed_ops.py src/native/ninetoothed/codegen.py tests/test_generate_ninetoothed_ops.py && python3 tests/test_generate_ninetoothed_ops.py && git diff --check
.
----------------------------------------------------------------------
Ran 1 test in 0.011s

OK

Remote Python checks:
$ ruff check scripts/generate_ninetoothed_ops.py src/native/ninetoothed/codegen.py tests/test_generate_ninetoothed_ops.py && ruff format --check scripts/generate_ninetoothed_ops.py src/native/ninetoothed/codegen.py tests/test_generate_ninetoothed_ops.py
All checks passed!
3 files already formatted

Remote NVIDIA build:
- `ssh nvidia`, `infiniops-ci/nvidia:latest`.
- `ntops` installed with `python3 -m pip install --no-deps git+https://github.com/InfiniTensor/ntops.git` after installing `ninetoothed==0.25.0`, `triton==3.7.0`, and `numpy>=1.26.4,<2` in the validation container.
- Built `ops` target with `WITH_NVIDIA=ON`, `WITH_NINETOOTHED=ON`, and `GENERATE_PYTHON_BINDINGS=ON`.
- Ran `ops.RmsNorm.active_implementation_indices('cuda')`: `[0, 9]`.
- Ran `ops.rms_norm(..., implementation_index=9)` on shapes `(13, 4)` and `(2, 3, 17)`, dtype `float32`.
- Compared against `torch.nn.functional.rms_norm`: max errors `2.384185791015625e-07` and `2.384185791015625e-07`, `torch.allclose(...)=True` for both.

Benchmark / Performance Impact

N/A. This PR wires a generated backend path and only includes correctness smoke tests; no benchmark claim is made.

Notes for Reviewers

scripts/generate_ninetoothed_ops.py is now only a build entrypoint; operator-specific generation lives in src/native/ninetoothed/codegen.py.
ntops.kernels.rms_norm.premake is used directly for the arrangement/application. The generator no longer replaces premake tensors with concrete-shape tensors; it generates by supported ranks via INFINIOPS_NINETOOTHED_RMS_NORM_NDIMS.
NineToothed uses NineToothedTensor even for scalar eps and num_normalized_elements; the slot 9 adapter wraps those scalars accordingly.
Common NineToothed tensor/dtype/size/scalar adapters live in src/native/ninetoothed/tensor.h; the RmsNorm header only keeps the 1D-weight expansion and generated launcher call.
ntops is expected to be importable in the selected Python environment. A local ninetoothed checkout is still optionally supported through NINETOOTHED_SOURCE_DIR for development.
Full default pytest and all-platform CI were not run in this iteration.

Checklist

Title, Branch, and Commits

PR title follows Conventional Commits (e.g. feat(nvidia): ..., fix(cuda/gemm): ...).
Branch name follows <type>/xxx-yyyy-zzzz where <type> matches the PR title's Conventional Commits type and words are joined with hyphens (see CONTRIBUTING.md §Branches).
Each commit message follows Conventional Commits.
Small PR is a single squashable commit; or, for a large PR, every commit is meaningful, well-formed, and independently reviewable (see CONTRIBUTING.md §Pull Requests).
No stray merge commits from master - the branch is rebased cleanly on top of the current master.
No fixup! / squash! / wip commits remain.

Scope and Design

Changes are minimal - nothing unrelated to the stated motivation was added (CONTRIBUTING.md §Code/General).
No dead code, commented-out blocks, debug prints, printf/std::cout/print(...) left behind, or TODO without an owner and issue link.
No unrelated formatting churn that would obscure the diff.
Public API changes (if any) are intentional, documented, and reflected in affected callers/tests.

General Code Hygiene (applies to all languages)

The code is self-explanatory; comments were added only where the why is non-obvious (CONTRIBUTING.md §Code/General).
Every modified or added file ends with a single trailing newline (CONTRIBUTING.md §Code/General).
No trailing whitespace, tab/space mixing, or stray BOMs.
Identifiers in comments and error messages are wrapped in backticks (e.g. the `seqlens_k` tensor) (CONTRIBUTING.md §Code/General).
All comments and error messages are in English (CONTRIBUTING.md §Code/General).
Comments and error messages are complete sentences - capitalized first letter, terminal punctuation - unless the language/framework convention says otherwise (CONTRIBUTING.md §Code/General; §Python).

C++ Specific (if C++ files changed)

Python Specific (if Python files changed)

Code is PEP 8 compliant; ruff check passes cleanly on CI (see .github/workflows/ruff.yml). Verified in infiniops-ci/nvidia:latest.
ruff format --check passes cleanly. Verified in infiniops-ci/nvidia:latest.
Comments are complete English sentences, starting with a capital letter and ending with punctuation; Markdown backticks are used for code references (CONTRIBUTING.md §Python).
Framework-specific conventions (e.g. lowercase pytest.skip messages without terminal period) are honored where applicable (CONTRIBUTING.md §Python).
No blank line between the function signature and the body when there is no docstring or comment (CONTRIBUTING.md §Python).
A blank line is present before and after if, for, and similar control-flow statements (CONTRIBUTING.md §Python).
A blank line appears before each return, except when it directly follows a control-flow statement (CONTRIBUTING.md §Python).
N/A - No docstrings were added.
Type hints are consistent with the surrounding script style.

Testing

pytest was run locally on every supported platform that this PR can affect, and the results are recorded in the "Test Results" table above (CONTRIBUTING.md §Pull Requests). Only targeted local and NVIDIA smoke checks were run.
For any platform that could not be tested, an explicit reason is given in the table and a reviewer with access has been tagged.
New functionality has matching tests under tests/ following tests/test_add.py / tests/test_gemm.py patterns (CONTRIBUTING.md §Adding an Operator).
N/A - The new generator unit test does not use dependent pytest.mark.parametrize parameters.
N/A - The new generator unit test is a unittest test for code generation.
N/A - The new generator unit test does not use default dtype/device parameterization.
N/A - No known parallel flakiness.
N/A - This is a feature PR, not a bug fix.

Build, CI, and Tooling

The project builds cleanly from a fresh directory with pip install .[dev] on at least one affected platform. Not run; targeted CMake/Ninja build was run in the NVIDIA CI image.
compile_commands.json still regenerates (CMake option CMAKE_EXPORT_COMPILE_COMMANDS=ON in pyproject.toml - required by the code-lint skill and clang-tidy -p).
N/A - No new backend/device was added.
Only one CUDA-like GPU backend is selectable at a time - the existing mutual-exclusion check in CMakeLists.txt is not broken.
Both CI workflows (clang-format.yml, ruff.yml) are green locally (or expected to be green on CI). ruff was verified in the NVIDIA CI container; clang-format was unavailable locally and in that container.
No new runtime dependency was added without updating pyproject.toml's [project.optional-dependencies] (or justified in the PR description). ntops is required only when explicitly enabling WITH_NINETOOTHED codegen.

Documentation

README.md, CONTRIBUTING.md, or inline docs updated when behavior, build flags, or developer workflow changed. Build flags are introduced in CMake only; no user-facing docs were added in this PR.
New operators, new dispatch helpers, or new public utilities are documented (docstring, header comment, or an addition to CONTRIBUTING.md §Some Code Explanations).
N/A - No user-visible breaking change.

Security and Safety

No secrets, access tokens, internal URLs, customer data, or personal hardware identifiers have been committed.
Third-party code is license-compatible and attributed.
No unsafe pointer arithmetic, uninitialized reads, or missing bounds checks were introduced.

voltjia · 2026-05-20T02:08:10Z

 option(WITH_TORCH "Enable PyTorch C++ backend" OFF)

+option(WITH_NINETOOTHED "Enable NineToothed-generated NVIDIA kernels" OFF)
+set(NINETOOTHED_PYTHON_EXECUTABLE "" CACHE FILEPATH "Python executable used to run ninetoothed code generation")


这部分主要是用来写 option 的，请把下面这堆 set 给挪到一个专门的 section。

已改。WITH_NINETOOTHED 仍然放在 option 区，下面这些 cache 变量已经挪到单独的 NineToothed code generation configuration section 里。

voltjia · 2026-05-20T04:49:16Z

+_SUPPORTED_OPS = ("rms_norm",)
+
+
+def _import_ninetoothed(source_dir):


为啥不直接是 import ninetoothed？且请不要使用缩写，直接使用全称 ninetoothed。

已改。现在实现移到了 src/native/ninetoothed/codegen.py，_import_ninetoothed 只在可选 source dir 需要时调整 sys.path，随后直接 import ninetoothed，变量名也不再用 nt 缩写。

voltjia · 2026-05-20T04:49:26Z

+    return nt
+
+
+def _import_ntops():


为啥不直接是 import ntops？

已改。_rms_norm_premake 里直接 import ntops，不再通过 importlib 或 _import_ntops 包一层。

voltjia · 2026-05-20T04:57:09Z

+
+_DEFAULT_DTYPES = ("float32", "float16", "bfloat16")
+
+_DEFAULT_RMS_NORM_SHAPES = (


具体算子相关的内容，应该放到 src 里，scripts 里面只放纯功能性工具或者构建相关脚本。

已改。scripts/generate_ninetoothed_ops.py 现在只作为构建入口，把 src 加到 sys.path 后委托给 native.ninetoothed.codegen.main()；具体算子和生成逻辑放到了 src/native/ninetoothed/codegen.py。

voltjia · 2026-05-20T04:57:35Z

+    return importlib.import_module("ntops")
+
+
+def _rms_norm_premake_rank2(dim0, dim1, dtype, block_size):


同上，这部分应当放在 src 中合适的地方，而不是在 scripts 下。以下同类问题不再赘述，但请一并修改。

已一并修改。RmsNorm 的 premake 包装、rank/dtype config、manifest 生成都挪到了 src/native/ninetoothed/codegen.py，scripts 下不再放算子细节。

voltjia · 2026-05-20T06:53:06Z

+    return arrangement, application, tensors
+
+
+def _rms_norm_premake_rank3(dim0, dim1, dim2, dtype, block_size):


这为啥要分 rank？不是只有 shape 不一样，那不是传个 shape 就行了嘛？

已改。之前按 rank 拆函数是为了让 ninetoothed.build 生成不同 launcher 参数；现在改成使用 ntops 自带的动态-rank premake，只按 ndim/dtype 生成同一个 infiniops_ninetoothed_rms_norm dispatcher，Python 侧不再拆 rank2/rank3 premake。

voltjia · 2026-05-20T10:05:35Z

+    return arrangement, application, tensors
+
+
+def _parse_shape(value):


这个函数是干嘛的？

已删除。现在不再按具体 shape 编译，也不需要解析 1x64 这类 shape 字符串；配置改为 INFINIOPS_NINETOOTHED_RMS_NORM_NDIMS / --rms-norm-ndims。

voltjia · 2026-05-20T10:06:14Z

+
+namespace detail {
+
+inline int NineToothedRmsNormDTypeIndex(DataType dtype) {


这种类似的 helper 不是应该是整个九齿 common 的嘛？不要放在 rms_norm 下面。

已改。DTypeIndex、SizeArg、FromTensor、FromScalar 都放到 src/native/ninetoothed/tensor.h 作为九齿 common helper；rms_norm/ninetoothed.h 只保留 ExpandedRmsNormWeight 这种算子特有适配和 generated launcher 调用。

voltjia · 2026-05-22T01:31:59Z

+_SUPPORTED_OPS = ("rms_norm",)
+
+
+def _import_ninetoothed(source_dir):


到底为啥需要这个 helper？去掉它，直接在 top-level import ninetoothed 就行。

voltjia · 2026-05-22T01:32:16Z


 option(WITH_TORCH "Enable PyTorch C++ backend" OFF)

+option(WITH_NINETOOTHED "Enable NineToothed-generated NVIDIA kernels" OFF)


此处不提及 NVIDIA，因为九齿的目标是跨平台，只是目前可能只暴露了 cuda caller，所以跟 PyTorch 的对齐即可。

voltjia · 2026-05-22T03:00:14Z

九齿的定位在算子库中应该跟 PyTorch 差不多，都可以接入到后端里，所以在文件结构上应该跟 PyTorch 平行，而不是放在 cuda 下，现在的九齿可能只有 cuda 这个 caller，但是生成的接口是一致的，只要后期增多了支持，就可以跨平台，跟 PyTorch 一样。

voltjia · 2026-05-22T23:25:50Z

+#ifndef INFINI_OPS_NVIDIA_RMS_NORM_NINETOOTHED_H_
+#define INFINI_OPS_NVIDIA_RMS_NORM_NINETOOTHED_H_
+
+#ifdef WITH_NINETOOTHED


就像评论 https://github.com/InfiniTensor/InfiniOps/pull/616/changes#r3285560677 所说，九齿应该是与 PyTorch 等对应的后端，所以是跟 torch 差不多的文件架构，而咱们算子库都是靠 build system 和脚本来确定最终产物，所以不要在 src 里面的文件使用 WITH_NINETOOTHED 这种类似的宏。事实上，在 C++ 中，我们应当尽量少地使用宏。

voltjia · 2026-05-22T23:26:33Z

+#include "rms_norm/infiniops_ninetoothed_rms_norm.h"
+
+#ifndef INFINIOPS_NINETOOTHED_BLOCK_SIZE
+#define INFINIOPS_NINETOOTHED_BLOCK_SIZE 256


C++ 中尽量不使用宏，尤其是这种可以被 constexpr 或者 const 替代的情况。

voltjia · 2026-05-22T23:41:22Z

+  };
+}
+
+template <typename NineToothedTensor, typename T>


这个模板参数的意义是什么？暂时是冗余的。如无必要，吴增实体。

voltjia · 2026-05-22T23:49:22Z

我这里有一套方案，可以看看是不是更好一些：我们直接提供一个 infini::ops::ninetoothed::Tensor 类，这个类里面去定义从 infini::ops::Tensor 或者 scalar 到它的 implicit conversion。这样用起来会不会更方便一些？也请考虑一些可能的风险，综合评判是否这么做。

voltjia · 2026-05-22T23:52:09Z

-
-    set(${out_var} "" PARENT_SCOPE)
-endfunction()
+# NineToothed code generation configuration.


请在后面创建一个关于 WITH_NINETOOTHED 的 if 吧，把这些 set 放到这个分支里吧。

voltjia · 2026-05-22T23:52:57Z

+import sys
+import tempfile
+import types
+import unittest


为什么引入了 unittest？请统一使用 pytest，与其他测试保持一致。

voltjia · 2026-05-22T23:55:57Z

+
+namespace infini::ops {
+
+namespace detail {


我看此处 detail 内部的函数比较少，且每个函数的内容也很少，可以考虑直接放在 Operator<RmsNorm, Device::Type::kNvidia, 9>::operator() 里面，暂时不单独抽成独立的 helper 了。

voltjia commented May 20, 2026

View reviewed changes

feat(nvidia): add ntops rms norm backend

0ad2354

voltjia force-pushed the feat/nvidia-ntops-rms-norm branch from fa89de9 to 0ad2354 Compare May 20, 2026 11:23

voltjia commented May 22, 2026

View reviewed changes

		_SUPPORTED_OPS = ("rms_norm",)


		def _import_ninetoothed(source_dir):


		_DEFAULT_DTYPES = ("float32", "float16", "bfloat16")

		_DEFAULT_RMS_NORM_SHAPES = (

		return importlib.import_module("ntops")


		def _rms_norm_premake_rank2(dim0, dim1, dtype, block_size):

		return arrangement, application, tensors


		def _rms_norm_premake_rank3(dim0, dim1, dim2, dtype, block_size):

		return arrangement, application, tensors


		def _parse_shape(value):


		namespace detail {

		inline int NineToothedRmsNormDTypeIndex(DataType dtype) {


		option(WITH_TORCH "Enable PyTorch C++ backend" OFF)

		option(WITH_NINETOOTHED "Enable NineToothed-generated NVIDIA kernels" OFF)

Conversation

voltjia commented May 20, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Motivation

Type of Change

Platforms Affected

Test Results on Supported Platforms

Benchmark / Performance Impact

Notes for Reviewers

Checklist

Title, Branch, and Commits

Scope and Design

General Code Hygiene (applies to all languages)

C++ Specific (if C++ files changed)

Python Specific (if Python files changed)

Testing

Build, CI, and Tooling

Documentation

Security and Safety

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

voltjia commented May 20, 2026 •

edited

Loading