feat: add GammaAddRmsNorm custom operator.

test/python_test/RegisterOps.cpp

@@ -601,6 +601,37 @@ std::tuple<at::Tensor, at::Tensor, at::Tensor> add_rms_norm_bias_impl_npu(
   return std::make_tuple(y, rstd, x);
 }
 
+std::tuple<at::Tensor, at::Tensor, at::Tensor>
+gamma_add_rms_norm_impl_npu(const at::Tensor& x1,
+                            const at::Tensor& x2,
+                            const at::Tensor& gamma,
+                            double epsilon,
+                            bool add_gamma_offset) {
+  auto sizes = x1.sizes().vec();
+  const int64_t dim_x = static_cast<int64_t>(sizes.size());
+  const int64_t dim_gamma = gamma.dim();
+  std::vector<int64_t> rstd_shape(sizes.begin(), sizes.end());
+  for (int64_t i = dim_x - dim_gamma; i < dim_x && i >= 0; ++i) {
+    rstd_shape[i] = 1;
+  }
+
+  at::Tensor y = at::empty(sizes, x1.options());
+  at::Tensor rstd =
+      at::empty(rstd_shape, x1.options().dtype(at::ScalarType::Float));
+  at::Tensor x = at::empty(sizes, x1.options());
+
+  EXEC_NPU_CMD(aclnnGammaAddRmsNorm,
+               x1,
+               x2,
+               gamma,
+               epsilon,
+               add_gamma_offset,
+               y,
+               rstd,
+               x);
+  return std::make_tuple(y, rstd, x);
+}
+
 std::tuple<at::Tensor, at::Tensor, at::Tensor, at::Tensor>
 moe_init_routing_custom_impl_npu(const at::Tensor& x,
                                  const at::Tensor& expert_idx,
@@ -1171,6 +1202,9 @@ PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
   m.def("scatter_nd_update_v2", &scatter_nd_update_v2_impl_npu, "scatter_nd_update_v2");
   m.def("hc_post", &hc_post_impl_npu, "hc_post");
   m.def("add_rms_norm_bias", &add_rms_norm_bias_impl_npu, "add_rms_norm_bias");
+  m.def("gamma_add_rms_norm",
+        &gamma_add_rms_norm_impl_npu,
+        "gamma_add_rms_norm");
   m.def("moe_init_routing_custom", &moe_init_routing_custom_impl_npu, "moe_init_routing_custom");
   m.def("moe_init_routing_v3", &moe_init_routing_v3_impl_npu, "moe_init_routing_v3");
   m.def("hc_pre_sinkhorn", &hc_pre_sinkhorn_impl_npu, "hc_pre_sinkhorn");

test/python_test/custom_ops.py

@@ -181,6 +181,12 @@ def add_rms_norm_bias_npu(x1, x2, gamma, beta=None, eps=1e-6):
     return custom_ops_lib.add_rms_norm_bias(x1, x2, gamma, beta, eps)
 
 
+def gamma_add_rms_norm_npu(x1, x2, gamma, eps=1e-6,
+                           add_gamma_offset=False):
+    return custom_ops_lib.gamma_add_rms_norm(
+        x1, x2, gamma, eps, add_gamma_offset)
+
+
 # hc_pre_sinkhorn (per-token: pre/post gating + sinkhorn-normalized comb_frag)
 # mixes last dim = 2*hc_mult + hc_mult^2 ([pre | post | comb]).
 # returns (y[bs,d] bf16, post[bs,hc_mult] fp32, comb_frag[bs,hc_mult,hc_mult] fp32)

test/python_test/test_gamma_add_rms_norm.py

@@ -0,0 +1,88 @@
+#!/usr/bin/env python3
+# Copyright 2026 The xLLM Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+import pytest
+import torch
+
+
+torch_npu = pytest.importorskip("torch_npu")
+custom_ops = pytest.importorskip("custom_ops")
+
+
+def gamma_add_rms_norm_golden(x1, x2, gamma, eps, add_gamma_offset):
+    dtype = x1.dtype
+    if dtype == torch.float16:
+        x = x1 + x2
+    else:
+        x = (x1.float() + x2.float()).to(dtype)
+
+    adjusted_gamma = gamma
+    if add_gamma_offset:
+        adjusted_gamma = (gamma + 1).to(dtype)
+
+    x_float = x.float()
+    variance = (x_float * x_float).mean(
+        dim=tuple(range(x.dim() - gamma.dim(), x.dim())), keepdim=True)
+    rstd = torch.rsqrt(variance + eps)
+    normalized = (x_float * rstd).to(dtype)
+    y = (normalized * adjusted_gamma).to(dtype)
+    return y, rstd.float(), x
+
+
+# (x shape, gamma shape) covers decode, prefill, and multi-dimensional gamma.
+CASES = [
+    ((1, 1024), (1024,)),
+    ((4, 2048), (2048,)),
+    ((16, 5120), (5120,)),
+    ((2, 3, 256), (3, 256)),
+]
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
+@pytest.mark.parametrize("add_gamma_offset", [False, True])
+@pytest.mark.parametrize("x_shape,gamma_shape", CASES)
+def test_gamma_add_rms_norm(x_shape, gamma_shape, add_gamma_offset, dtype):
+    torch.manual_seed(20260716)
+    eps = 1e-6
+    x1 = (torch.rand(x_shape, dtype=torch.float32) - 0.5).to(dtype)
+    x2 = (torch.rand(x_shape, dtype=torch.float32) - 0.5).to(dtype)
+    gamma = (torch.rand(gamma_shape, dtype=torch.float32) - 0.5).to(dtype)
+
+    y_ref, rstd_ref, x_ref = gamma_add_rms_norm_golden(
+        x1, x2, gamma, eps, add_gamma_offset)
+    y_npu, rstd_npu, x_npu = custom_ops.gamma_add_rms_norm_npu(
+        x1.npu(),
+        x2.npu(),
+        gamma.npu(),
+        eps,
+        add_gamma_offset,
+    )
+    torch.npu.synchronize()
+
+    if dtype == torch.float16:
+        atol = rtol = 1e-3
+    elif dtype == torch.bfloat16:
+        atol = rtol = 5e-3
+    else:
+        atol = rtol = 1e-5
+
+    assert tuple(rstd_npu.shape) == tuple(rstd_ref.shape)
+    torch.testing.assert_close(
+        x_npu.cpu().float(), x_ref.float(), atol=atol, rtol=rtol)
+    torch.testing.assert_close(
+        rstd_npu.cpu().float(), rstd_ref, atol=1e-5, rtol=1e-5)
+    torch.testing.assert_close(
+        y_npu.cpu().float(), y_ref.float(), atol=atol, rtol=rtol)

xllm_ops/build_aclnn.sh

@@ -110,6 +110,7 @@ elif [[ "$SOC_VERSION" =~ ^(ascend)?910b ]]; then
         "moe_init_routing_custom"
         "moe_gating_top_k_hash"
         "add_rms_norm_bias"
+        "gamma_add_rms_norm"
         "lightning_indexer_quant"
         "compressor"
         "quant_lightning_indexer"  ## 已在 CANN 中内置，见 opp/built-in/op_impl/ai_core/tbe/impl/ops_transformer/ascendc/quant_lightning_indexer
@@ -218,6 +219,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend910_93 ]]; then
         "moe_init_routing_custom"
         "moe_gating_top_k_hash"
         "add_rms_norm_bias"
+        "gamma_add_rms_norm"
         "lightning_indexer_quant"
         "lightning_indexer_quant_metadata"
         "compressor"
@@ -290,6 +292,7 @@ elif [[ "$SOC_VERSION" =~ ^ascend950 ]]; then
         "moe_init_routing_custom"
         "moe_gating_top_k_hash"
         "add_rms_norm_bias"
+        "gamma_add_rms_norm"
         "lightning_indexer_quant" 
         "lightning_indexer_quant_metadata"
         "compressor"

xllm_ops/gamma_add_rms_norm/CMakeLists.txt

@@ -0,0 +1,8 @@
+# Copyright 2026 The xLLM Authors. All Rights Reserved.
+
+file(GLOB CURRENT_DIRS RELATIVE ${CMAKE_CURRENT_SOURCE_DIR} ${CMAKE_CURRENT_SOURCE_DIR}/*)
+foreach(SUB_DIR ${CURRENT_DIRS})
+    if(EXISTS "${CMAKE_CURRENT_SOURCE_DIR}/${SUB_DIR}/CMakeLists.txt")
+        add_subdirectory(${SUB_DIR})
+    endif()
+endforeach()