From ab7897b150faee5e0b877717fb4284880507182b Mon Sep 17 00:00:00 2001
From: WangHuanjun <wanghuanjun2113@users.noreply.github.com>
Date: Mon, 13 Jul 2026 18:49:48 +0800
Subject: [PATCH] perf: add native split gated rms norm path.

- add an opt-in NPU path using native RMSNorm, SiLU, and in-place multiplication
- preserve the fused implementation as the default and fallback path
- limit the optimization to large two-dimensional inputs with matching norm dimensions
---
 xllm/core/layers/common/rms_norm_gated.cpp | 35 ++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/xllm/core/layers/common/rms_norm_gated.cpp b/xllm/core/layers/common/rms_norm_gated.cpp
index b6027f9a75..7ebdc15d1f 100644
--- a/xllm/core/layers/common/rms_norm_gated.cpp
+++ b/xllm/core/layers/common/rms_norm_gated.cpp
@@ -15,14 +15,38 @@ limitations under the License.
 
 #include "rms_norm_gated.h"
 
+#if defined(USE_NPU)
+#include <ATen/ops/silu.h>
+
+#include <cstdlib>
+#include <cstring>
+#endif
+
 #include <glog/logging.h>
 
 #include "framework/state_dict/utils.h"
+#if defined(USE_NPU)
+#include "xllm/core/kernels/npu/npu_ops_api.h"
+#endif
 #include "xllm/core/kernels/ops_api.h"
 
 namespace xllm {
 namespace layer {
 
+#if defined(USE_NPU)
+namespace {
+
+bool enable_native_split_gated_rms_norm() {
+  static const bool enabled = []() {
+    const char* value = std::getenv("XLLM_NATIVE_SPLIT_GATED_RMS_NORM");
+    return value != nullptr && std::strcmp(value, "1") == 0;
+  }();
+  return enabled;
+}
+
+}  // namespace
+#endif
+
 RmsNormGatedImpl::RmsNormGatedImpl(int64_t dim,
                                    double eps,
                                    const torch::TensorOptions& options)
@@ -33,6 +57,17 @@ RmsNormGatedImpl::RmsNormGatedImpl(int64_t dim,
 
 torch::Tensor RmsNormGatedImpl::forward(torch::Tensor& input,
                                         std::optional<torch::Tensor> gate) {
+#if defined(USE_NPU)
+  if (enable_native_split_gated_rms_norm() && gate.has_value() &&
+      input.dim() == 2 && input.size(0) >= 1024 && input.size(1) == norm_dim_) {
+    auto normalized =
+        xllm::kernel::npu::rms_norm(input, weight_, eps_, "rmsnorm");
+    auto gate_value = gate.value();
+    at::silu_(gate_value);
+    normalized.mul_(gate_value);
+    return normalized;
+  }
+#endif
   xllm::kernel::GatedLayerNormParams params;
   params.x = input;
   params.weight = weight_;