feat: integrate add_rms_norm interface for NPU backend.

Author: yingxudeng

xllm/core/kernels/npu/attention.cpp

@@ -33,8 +33,8 @@ void batch_prefill(const torch::Tensor& query,
                    const torch::Tensor& seq_len,
                    float scale,
                    torch::Tensor& output) {
-  auto num_heads = query.size(-2);
-  auto num_kv_heads = key.size(-2);
+  int64_t num_heads = query.size(-2);
+  int64_t num_kv_heads = key.size(-2);
   atb::_npu_flash_attention(
       query, key, value, mask, seq_len, scale, num_heads, num_kv_heads, output);
 }

xllm/core/kernels/npu/fused_layernorm.cpp

@@ -33,4 +33,12 @@ torch::Tensor rms_norm(const torch::Tensor& input,
   return normalized_input;
 }
 
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> add_rms_norm(
+    const torch::Tensor& x1,
+    const torch::Tensor& x2,
+    const torch::Tensor& gamma,
+    double epsilon) {
+  return at_npu::native::custom_ops::npu_add_rms_norm(x1, x2, gamma, epsilon);
+}
+
 }  // namespace xllm::kernel::npu

xllm/core/kernels/npu/npu_ops_api.h

@@ -17,6 +17,7 @@ limitations under the License.
 #include <torch/torch.h>
 
 #include <optional>
+#include <tuple>
 
 #include "custom_functions_npu/atb_common.h"
 
@@ -55,6 +56,12 @@ torch::Tensor rms_norm(const torch::Tensor& input,
                        double eps,
                        const std::string& mode);
 
+std::tuple<torch::Tensor, torch::Tensor, torch::Tensor> add_rms_norm(
+    const torch::Tensor& x1,
+    const torch::Tensor& x2,
+    const torch::Tensor& gamma,
+    double epsilon);
+
 void apply_rotary(torch::Tensor& q,
                   torch::Tensor& k,
                   const torch::Tensor& cos_sin_cache,

xllm/core/kernels/npu/rope.cpp

@@ -25,7 +25,7 @@ void apply_rotary(torch::Tensor& q,
                   const torch::Tensor& cos_sin_cache,
                   const torch::Tensor& positions) {
   auto cos_sin = cos_sin_cache.index_select(0, positions);
-  auto last_dim = cos_sin.size(-1);
+  int64_t last_dim = cos_sin.size(-1);
   auto cos_sin_vec = cos_sin.view({-1, 2, last_dim / 2})
                          .repeat({1, 1, 2})
                          .chunk(2, /*dim=*/-2);

xllm/core/kernels/ops_api.cpp

@@ -284,8 +284,14 @@ void fused_layernorm(FusedLayerNormParams& params) {
                        params.store_output_after_norm,
                        params.dynamic_quant);
 #elif defined(USE_NPU)
-  params.output =
-      npu::rms_norm(params.input, params.weight, params.eps, params.mode);
+  if (params.residual.has_value()) {
+    std::tie(params.output, std::ignore, params.residual_out) =
+        npu::add_rms_norm(
+            params.input, params.residual.value(), params.weight, params.eps);
+  } else {
+    params.output =
+        npu::rms_norm(params.input, params.weight, params.eps, params.mode);
+  }
 #elif defined(USE_CUDA)
   if (params.residual.has_value()) {
     cuda::fused_add_rms_norm(