Skip to content

Commit 9cd53bc

Browse files
voltjiavoltjia
committed
feat: add NineToothed fused add RMSNorm
1 parent 3ac6719 commit 9cd53bc

13 files changed

Lines changed: 241 additions & 22 deletions

File tree

‎.github/ci_config.yml‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ platforms:
1515
BASE_IMAGE: nvcr.io/nvidia/pytorch:24.10-py3
1616
SKIP_APT: "1"
1717
PIP_INDEX_URL: https://pypi.tuna.tsinghua.edu.cn/simple
18-
setup: pip install scikit-build-core pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_NVIDIA=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
18+
setup: pip install "scikit-build-core>=1.1.1" pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix /tmp/ninetoothed-src /tmp/ntops-src && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_NVIDIA=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && git clone ${NINETOOTHED_GIT_URL:-https://github.com/InfiniTensor/ninetoothed.git} /tmp/ninetoothed-src && git -C /tmp/ninetoothed-src checkout dfce2a5f0ec2a2a9428cb4cc897a344eb7a8f725 && pip install /tmp/ninetoothed-src && git clone --depth 1 ${NTOPS_GIT_URL:-https://github.com/InfiniTensor/ntops.git} /tmp/ntops-src && git -C /tmp/ntops-src checkout 82cba6cf2200919b5d12c86adbea60cbfb516040 && pip install /tmp/ntops-src && pip install "numpy<2" && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.WITH_NINETOOTHED=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
1919
jobs:
2020
gpu:
2121
type: unittest
@@ -51,7 +51,7 @@ platforms:
5151
- /lib/firmware:/lib/firmware
5252
- /usr/src:/usr/src
5353
- /lib/modules:/lib/modules
54-
setup: python -m pip install scikit-build-core packaging exceptiongroup typing-extensions pygments pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_ILUVATAR=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix -DCMAKE_CUDA_ARCHITECTURES=OFF -DCMAKE_CUDA_COMPILER=/usr/local/corex-4.3.0.20250624/bin/clang++ -DCUDAToolkit_ROOT=/usr/local/corex-4.3.0.20250624 && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && python -m pip install . --no-build-isolation --no-deps --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
54+
setup: python -m pip install "scikit-build-core>=1.1.1" packaging exceptiongroup typing-extensions pygments pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_ILUVATAR=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix -DCMAKE_CUDA_ARCHITECTURES=OFF -DCMAKE_CUDA_COMPILER=/usr/local/corex-4.3.0.20250624/bin/clang++ -DCUDAToolkit_ROOT=/usr/local/corex-4.3.0.20250624 && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && python -m pip install . --no-build-isolation --no-deps --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
5555
jobs:
5656
gpu:
5757
type: unittest
@@ -81,7 +81,7 @@ platforms:
8181
- "--privileged"
8282
- "--ulimit=memlock=-1"
8383
- "--ulimit=stack=67108864"
84-
setup: pip install scikit-build-core pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_METAX=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
84+
setup: pip install "scikit-build-core>=1.1.1" pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_METAX=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && MACA_PATH=/opt/maca SKBUILD_CMAKE_DEFINE="CMAKE_C_COMPILER=$PWD/scripts/mxcc_wrapper.sh;CMAKE_CXX_COMPILER=$PWD/scripts/mxcc_wrapper.sh;WITH_METAX=ON;GENERATE_PYTHON_BINDINGS=ON;INFINI_OPS_SMOKE_BUILD=ON;INFINI_RT_ROOT=/tmp/infinirt-prefix" pip install .[dev] --no-build-isolation
8585
jobs:
8686
gpu:
8787
type: unittest
@@ -109,7 +109,7 @@ platforms:
109109
PIP_INDEX_URL: https://pypi.org/simple
110110
docker_args:
111111
- "--privileged"
112-
setup: pip install scikit-build-core pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_MOORE=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
112+
setup: pip install "scikit-build-core>=1.1.1" pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_MOORE=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && SKBUILD_CMAKE_DEFINE="CMAKE_C_COMPILER=$PWD/scripts/mcc_wrapper.sh;CMAKE_CXX_COMPILER=$PWD/scripts/mcc_wrapper.sh;WITH_MOORE=ON;GENERATE_PYTHON_BINDINGS=ON;INFINI_OPS_SMOKE_BUILD=ON;INFINI_RT_ROOT=/tmp/infinirt-prefix" pip install .[dev] --no-build-isolation
113113
jobs:
114114
gpu:
115115
type: unittest
@@ -124,7 +124,7 @@ platforms:
124124
junit_path: test-results.xml
125125
stages:
126126
- name: test
127-
run: pytest tests/ -m smoke -n 4 -v --tb=short --junitxml=/workspace/results/test-results.xml
127+
run: LD_LIBRARY_PATH=/usr/local/musa/lib:$LD_LIBRARY_PATH pytest tests/ -m smoke -n 4 -v --tb=short --junitxml=/workspace/results/test-results.xml
128128

129129
cambricon:
130130
runner_label: Cambricon
@@ -136,7 +136,7 @@ platforms:
136136
PIP_INDEX_URL: https://pypi.org/simple
137137
docker_args:
138138
- "--privileged"
139-
setup: pip install scikit-build-core pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_CAMBRICON=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
139+
setup: pip install "scikit-build-core>=1.1.1" pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_CAMBRICON=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
140140
jobs:
141141
gpu:
142142
type: unittest
@@ -172,7 +172,7 @@ platforms:
172172
- "--group-add=video"
173173
volumes:
174174
- /opt/hyhal:/opt/hyhal:ro
175-
setup: pip install scikit-build-core pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_HYGON=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
175+
setup: pip install "scikit-build-core>=1.1.1" pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_HYGON=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
176176
jobs:
177177
gpu:
178178
type: unittest
@@ -207,7 +207,7 @@ platforms:
207207
- /usr/local/bin/npu-smi:/usr/local/bin/npu-smi:ro
208208
env:
209209
ASCEND_HOME_PATH: /usr/local/Ascend/ascend-toolkit/latest
210-
setup: pip install scikit-build-core pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://gh-proxy.com/https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_ASCEND=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
210+
setup: pip install "scikit-build-core>=1.1.1" pybind11 libclang pyyaml && rm -rf /tmp/infinirt-src /tmp/infinirt-build /tmp/infinirt-prefix && git clone --depth 1 ${INFINI_RT_GIT_URL:-https://github.com/InfiniTensor/InfiniRT.git} /tmp/infinirt-src && cmake -S /tmp/infinirt-src -B /tmp/infinirt-build -DWITH_CPU=ON -DWITH_ASCEND=ON -DCMAKE_BUILD_TYPE=Release -DCMAKE_INSTALL_PREFIX=/tmp/infinirt-prefix && cmake --build /tmp/infinirt-build -j$(nproc) && cmake --install /tmp/infinirt-build && pip install .[dev] --no-build-isolation --config-settings=cmake.define.INFINI_OPS_SMOKE_BUILD=ON --config-settings=cmake.define.INFINI_RT_ROOT=/tmp/infinirt-prefix
211211
jobs:
212212
npu:
213213
type: unittest

‎README.md‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,7 @@ InfiniOps extension so `import infini.ops` can load its runtime dependency.
4848
| `-DWITH_MOORE=[ON\|OFF]` | Compile the Moore implementation | OFF |
4949
| `-DWITH_CAMBRICON=[ON\|OFF]` | Compile the Cambricon implementation | OFF |
5050
| `-DWITH_ASCEND=[ON\|OFF]` | Compile the Ascend implementation | OFF |
51+
| `-DWITH_NINETOOTHED=[ON\|OFF]` | Compile NineToothed-generated NVIDIA kernels; also requires `-DWITH_NVIDIA=ON` and Python packages `ninetoothed` and `ntops` | OFF |
5152
| `-DWITH_TORCH=[ON\|OFF]` | Compile generated PyTorch ATen-backed operators | OFF |
5253
| `-DBUILD_NVIDIA_W4_MARLIN=[ON\|OFF]` | Build AWQ/GPTQ Marlin providers; requires CUDA >= 12.3 and C++20 | OFF |
5354
| `-DAUTO_DETECT_DEVICES=[ON\|OFF]` | Auto-detect available platforms | ON |

‎scripts/generate_public_operator_header.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -105,7 +105,7 @@
105105
r"const\s+Config\s*&\s*config\s*,\s*"
106106
r"Device::Type\s+dispatch_device_type\s*,\s*"
107107
r"bool\s+strict\s*,\s*Selected\s*&&\s*selected\s*,\s*"
108-
r"Args\s*&&\s*\.\.\.\s*args\s*\)\s*",
108+
r"Args\s*&&\s*\.\.\.\s*(?:args|call_args)\s*\)\s*",
109109
re.MULTILINE,
110110
),
111111
),

‎src/CMakeLists.txt‎

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -642,6 +642,10 @@ if(WITH_NVIDIA)
642642
list(APPEND _infini_ops_smoke_ops cutlass_scaled_mm moe_sum)
643643
endif()
644644

645+
if(WITH_NVIDIA AND WITH_NINETOOTHED)
646+
list(APPEND _infini_ops_smoke_ops fused_add_rms_norm)
647+
endif()
648+
645649
if(((WITH_NVIDIA OR WITH_MOORE) AND WITH_TORCH) OR WITH_ILUVATAR)
646650
list(APPEND _infini_ops_smoke_ops flash_attn_varlen_func)
647651
endif()
@@ -1434,6 +1438,8 @@ if(GENERATE_OPERATOR_CALL_INSTANTIATIONS)
14341438
"${PROJECT_SOURCE_DIR}/generated/bindings/generated_dispatch_*.cc")
14351439

14361440
if(OPERATOR_DISPATCH_SOURCES)
1441+
configure_file("${CMAKE_CURRENT_SOURCE_DIR}/operator.h"
1442+
"${PROJECT_SOURCE_DIR}/generated/bindings/operator.h" COPYONLY)
14371443
target_include_directories(infiniops PRIVATE "${PROJECT_SOURCE_DIR}")
14381444
if(WITH_NVIDIA OR WITH_HYGON OR WITH_THEAD)
14391445
set_source_files_properties(${OPERATOR_DISPATCH_SOURCES}

‎src/base/cast.h‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -7,6 +7,8 @@ namespace infini::ops {
77

88
class Cast : public Operator<Cast> {
99
public:
10+
static constexpr bool kAutomaticDefault = true;
11+
1012
Cast(const Tensor input, Tensor out)
1113
: ndim_{out.ndim()},
1214
output_size_{out.numel()},

‎src/native/cuda/nvidia/ops/quick_gelu/kernel.cu‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -11,16 +11,16 @@
1111

1212
namespace infini::ops::nvidia_quick_gelu_detail {
1313
namespace {
14-
// The three branches preserve Core 0.2.9 quickgelu/cuda/kernel.cuh semantics.
15-
// In particular FP16 must not be promoted to the BF16 FP32-intermediate path.
14+
// CUDA 12.6 exposes colliding FP16/BF16 overloads for these intrinsic names.
15+
// Keep FP16 on the same FP32-intermediate path used by BF16.
1616
template <typename T>
1717
__device__ __forceinline__ T Apply(T x) {
1818
constexpr float alpha = 1.702f;
1919
if constexpr (std::is_same_v<T, __half>) {
20-
const __half ax = __hmul(__float2half(alpha), x);
21-
const __half denominator = __hadd(__float2half(1.0f), hexp(__hneg(ax)));
22-
const __half sigmoid = hrcp(denominator);
23-
return __hmul(x, sigmoid);
20+
const float xf = __half2float(x);
21+
const float ax = alpha * xf;
22+
const float sigmoid = 1.0f / (1.0f + __expf(-ax));
23+
return __float2half(xf * sigmoid);
2424
} else if constexpr (std::is_same_v<T, __nv_bfloat16>) {
2525
const float xf = __bfloat162float(x);
2626
const float ax = alpha * xf;

‎src/ninetoothed/ops/fused_add_rms_norm/build.py‎

Lines changed: 89 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,89 @@
1+
import itertools
2+
3+
import ninetoothed
4+
import ninetoothed.language as ntl
5+
6+
_BLOCK_SIZE = 1024
7+
8+
_DTYPES = ("float32", "float16", "bfloat16")
9+
10+
_DEFAULT_NDIMS = (2, 3, 4)
11+
12+
_CONFIGS = tuple(
13+
(
14+
(),
15+
{
16+
"ndim": ndim,
17+
"dtype": dtype,
18+
"block_size": _BLOCK_SIZE,
19+
},
20+
{},
21+
)
22+
for ndim, dtype in itertools.product(
23+
_DEFAULT_NDIMS,
24+
(getattr(ninetoothed, name) for name in _DTYPES),
25+
)
26+
)
27+
28+
29+
def premake(ndim, dtype=None, block_size=None):
30+
if block_size is None:
31+
raise ValueError("FusedAddRmsNorm requires an explicit block size")
32+
33+
tile_shape = (1,) * (ndim - 1) + (block_size,)
34+
35+
def arrangement_(input, residual, weight, has_weight, eps, num_normalized_elements):
36+
def arrange(tensor):
37+
return tensor.tile(tile_shape)
38+
39+
return (
40+
arrange(input),
41+
arrange(residual),
42+
arrange(weight),
43+
has_weight,
44+
eps,
45+
num_normalized_elements,
46+
)
47+
48+
def application(input, residual, weight, has_weight, eps, num_normalized_elements):
49+
input_f32 = ntl.cast(input, ntl.float32)
50+
residual_f32 = ntl.cast(residual, ntl.float32)
51+
summed = input_f32 + residual_f32
52+
rounded_sum = ntl.cast(summed, residual.dtype)
53+
normalized = ntl.cast(rounded_sum, ntl.float32)
54+
scale = ntl.rsqrt(
55+
ntl.sum(normalized * normalized) / num_normalized_elements + eps
56+
)
57+
residual = rounded_sum
58+
weight_multiplier = ntl.cast(has_weight, ntl.float32)
59+
weight_multiplier = weight_multiplier * ntl.cast(weight, ntl.float32) + (
60+
1 - weight_multiplier
61+
)
62+
input = ntl.cast(normalized * scale * weight_multiplier, input.dtype)
63+
64+
return (
65+
arrangement_,
66+
application,
67+
(
68+
ninetoothed.Tensor(ndim, other=0, dtype=dtype),
69+
ninetoothed.Tensor(ndim, other=0, dtype=dtype),
70+
ninetoothed.Tensor(ndim, dtype=dtype),
71+
ninetoothed.Tensor(0, dtype=ninetoothed.int32),
72+
ninetoothed.Tensor(0, dtype=ninetoothed.float64),
73+
ninetoothed.Tensor(0, dtype=ninetoothed.int64),
74+
),
75+
)
76+
77+
78+
def build(output_dir):
79+
variant_dir = output_dir / "fused_add_rms_norm"
80+
variant_dir.mkdir(parents=True, exist_ok=True)
81+
ninetoothed.build(
82+
premake,
83+
_CONFIGS,
84+
meta_parameters=("block_size",),
85+
caller="cuda",
86+
kernel_name="infini_ops_ninetoothed_fused_add_rms_norm",
87+
output_dir=variant_dir,
88+
lazy=False,
89+
)

0 commit comments

Comments
 (0)