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 + +#include +#include +#endif + #include #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 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_;