From d8c2c613131825b9b95dd1495a29f64e5c48473f Mon Sep 17 00:00:00 2001 From: Fengfengst123 <1322434659@qq.com> Date: Tue, 14 Jul 2026 22:25:46 +0800 Subject: [PATCH 1/2] feat(npu): route gated norm to aclnn layer norm fwd --- third_party/xllm_ops | 2 +- xllm/core/kernels/npu/xllm_ops/CMakeLists.txt | 1 + .../kernels/npu/xllm_ops/layer_norm_fwd.cpp | 116 ++++++++++++++++++ xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h | 10 ++ xllm/core/kernels/ops_api.cpp | 30 +++++ 5 files changed, 158 insertions(+), 1 deletion(-) create mode 100644 xllm/core/kernels/npu/xllm_ops/layer_norm_fwd.cpp diff --git a/third_party/xllm_ops b/third_party/xllm_ops index 9e436ee34a..4493779f11 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit 9e436ee34a69d25a9806d537c2a8de1f525b1331 +Subproject commit 4493779f11b5dc6fe431c63c1b895d6721831e33 diff --git a/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt b/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt index 459f53b85a..6cf0dacd1c 100644 --- a/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt +++ b/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt @@ -35,6 +35,7 @@ cc_library( rec_constrained_topk_fused.cpp beam_search_rec_final.cpp npu_mega_chunk_gdn.cpp + layer_norm_fwd.cpp INCLUDES ${PROJECT_SOURCE_DIR}/third_party/xllm_ops/build/autogen DEPS diff --git a/xllm/core/kernels/npu/xllm_ops/layer_norm_fwd.cpp b/xllm/core/kernels/npu/xllm_ops/layer_norm_fwd.cpp new file mode 100644 index 0000000000..029f0737ce --- /dev/null +++ b/xllm/core/kernels/npu/xllm_ops/layer_norm_fwd.cpp @@ -0,0 +1,116 @@ +/* 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 + + https://github.com/jd-opensource/xllm/blob/main/LICENSE + +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. +==============================================================================*/ + +#include "core/kernels/npu/aclnn/pytorch_npu_helper.hpp" +#include "xllm_ops_api.h" + +namespace xllm::kernel::npu { +namespace { + +bool is_supported_layer_norm_dtype(at::ScalarType dtype) { + return dtype == at::kHalf || dtype == at::kBFloat16 || dtype == at::kFloat; +} + +torch::Tensor make_contiguous_2d(const torch::Tensor& tensor) { + return tensor.reshape({-1, tensor.size(-1)}).contiguous(); +} + +} // namespace + +torch::Tensor layer_norm_fwd_aclnn(const torch::Tensor& x, + const torch::Tensor& weight, + const torch::Tensor& bias, + double eps, + const std::optional& z, + int64_t group_size, + bool norm_before_gate, + bool is_rms_norm) { + TORCH_CHECK(x.defined(), "layer_norm_fwd_aclnn: x must be defined"); + TORCH_CHECK(weight.defined(), + "layer_norm_fwd_aclnn: weight must be defined"); + TORCH_CHECK(x.dim() >= 1, + "layer_norm_fwd_aclnn: x must have at least 1 dim"); + TORCH_CHECK(is_supported_layer_norm_dtype(x.scalar_type()), + "layer_norm_fwd_aclnn: x dtype must be fp16, bf16 or fp32, got ", + x.scalar_type()); + TORCH_CHECK(weight.dim() == 1, + "layer_norm_fwd_aclnn: weight must be 1D"); + + const int64_t full_n = x.size(-1); + if (group_size < 0) { + group_size = full_n; + } + TORCH_CHECK(group_size > 0, + "layer_norm_fwd_aclnn: group_size must be positive"); + TORCH_CHECK(full_n % group_size == 0, + "layer_norm_fwd_aclnn: last dim ", + full_n, + " must be divisible by group_size ", + group_size); + TORCH_CHECK(weight.numel() == full_n, + "layer_norm_fwd_aclnn: weight numel must equal x last dim"); + + if (bias.defined()) { + TORCH_CHECK(bias.dim() == 1, + "layer_norm_fwd_aclnn: bias must be 1D"); + TORCH_CHECK(bias.numel() == full_n, + "layer_norm_fwd_aclnn: bias numel must equal x last dim"); + } + if (z.has_value() && z->defined()) { + TORCH_CHECK(z->sizes() == x.sizes(), + "layer_norm_fwd_aclnn: z shape must match x"); + } + + torch::Tensor x_2d = make_contiguous_2d(x); + torch::Tensor weight_contiguous = weight.contiguous(); + + c10::optional bias_contiguous = c10::nullopt; + if (bias.defined()) { + bias_contiguous = bias.contiguous(); + } + + c10::optional z_2d = c10::nullopt; + if (z.has_value() && z->defined()) { + z_2d = make_contiguous_2d(*z); + } + + torch::Tensor y_2d = torch::empty_like(x_2d); + const int64_t m = x_2d.size(0); + const int64_t group_count = full_n / group_size; + torch::Tensor mean = + is_rms_norm + ? torch::empty({0}, x.options().dtype(at::kFloat)) + : torch::empty({group_count * m}, x.options().dtype(at::kFloat)); + torch::Tensor rstd = + torch::empty({group_count * m}, x.options().dtype(at::kFloat)); + + const float eps_f = static_cast(eps); + EXEC_NPU_CMD(aclnnLayerNormFwd, + x_2d, + weight_contiguous, + bias_contiguous, + z_2d, + eps_f, + group_size, + norm_before_gate, + is_rms_norm, + y_2d, + mean, + rstd); + + return y_2d.reshape(x.sizes()); +} + +} // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h b/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h index 52e076cc4a..1a4da99f23 100644 --- a/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h +++ b/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h @@ -383,4 +383,14 @@ std::pair npu_mega_chunk_gdn( const std::optional& cu_seqlens = std::nullopt, c10::ArrayRef q_seq_lens = {}, bool use_qk_l2norm_in_kernel = false); + +torch::Tensor layer_norm_fwd_aclnn( + const torch::Tensor& x, + const torch::Tensor& weight, + const torch::Tensor& bias, + double eps, + const std::optional& z = std::nullopt, + int64_t group_size = -1, + bool norm_before_gate = true, + bool is_rms_norm = false); } // namespace xllm::kernel::npu diff --git a/xllm/core/kernels/ops_api.cpp b/xllm/core/kernels/ops_api.cpp index b38addd7df..8fd37578e8 100644 --- a/xllm/core/kernels/ops_api.cpp +++ b/xllm/core/kernels/ops_api.cpp @@ -37,6 +37,8 @@ limitations under the License. #include "dcu/hipblaslt_fp8_adapter.h" #endif +#include +#include #include #include "common/macros.h" @@ -49,6 +51,24 @@ namespace { bool is_supported_initial_state_dtype(torch::ScalarType dtype) { return dtype == torch::kBFloat16 || dtype == torch::kFloat32; } + +bool env_flag_enabled(const char* name, bool default_value) { + const char* value = std::getenv(name); + if (value == nullptr) { + return default_value; + } + if (std::strcmp(value, "1") == 0 || std::strcmp(value, "true") == 0 || + std::strcmp(value, "TRUE") == 0 || std::strcmp(value, "on") == 0 || + std::strcmp(value, "ON") == 0) { + return true; + } + if (std::strcmp(value, "0") == 0 || std::strcmp(value, "false") == 0 || + std::strcmp(value, "FALSE") == 0 || std::strcmp(value, "off") == 0 || + std::strcmp(value, "OFF") == 0) { + return false; + } + return default_value; +} #endif #if defined(USE_DCU) @@ -1616,6 +1636,16 @@ std::tuple moe_gating_top_k_hash( torch::Tensor gated_layer_norm(GatedLayerNormParams& params) { #if defined(USE_NPU) + if (env_flag_enabled("XLLM_USE_ACLNN_LAYER_NORM_FWD", true)) { + return npu::layer_norm_fwd_aclnn(params.x, + params.weight, + params.bias, + params.eps, + params.z, + params.group_size, + params.norm_before_gate, + params.is_rms_norm); + } return npu::layer_norm_fwd(params.x, params.weight, params.bias, From 1d248cff6e8cd0334b9b82ef2d487df0d36c9a25 Mon Sep 17 00:00:00 2001 From: Fengfengst123 <1322434659@qq.com> Date: Fri, 17 Jul 2026 11:53:32 +0800 Subject: [PATCH 2/2] refactor(npu): use aclnn layer norm fwd directly --- xllm/core/kernels/ops_api.cpp | 46 ++++++----------------------------- 1 file changed, 8 insertions(+), 38 deletions(-) diff --git a/xllm/core/kernels/ops_api.cpp b/xllm/core/kernels/ops_api.cpp index 8fd37578e8..24c7450e44 100644 --- a/xllm/core/kernels/ops_api.cpp +++ b/xllm/core/kernels/ops_api.cpp @@ -37,8 +37,6 @@ limitations under the License. #include "dcu/hipblaslt_fp8_adapter.h" #endif -#include -#include #include #include "common/macros.h" @@ -51,24 +49,6 @@ namespace { bool is_supported_initial_state_dtype(torch::ScalarType dtype) { return dtype == torch::kBFloat16 || dtype == torch::kFloat32; } - -bool env_flag_enabled(const char* name, bool default_value) { - const char* value = std::getenv(name); - if (value == nullptr) { - return default_value; - } - if (std::strcmp(value, "1") == 0 || std::strcmp(value, "true") == 0 || - std::strcmp(value, "TRUE") == 0 || std::strcmp(value, "on") == 0 || - std::strcmp(value, "ON") == 0) { - return true; - } - if (std::strcmp(value, "0") == 0 || std::strcmp(value, "false") == 0 || - std::strcmp(value, "FALSE") == 0 || std::strcmp(value, "off") == 0 || - std::strcmp(value, "OFF") == 0) { - return false; - } - return default_value; -} #endif #if defined(USE_DCU) @@ -1636,24 +1616,14 @@ std::tuple moe_gating_top_k_hash( torch::Tensor gated_layer_norm(GatedLayerNormParams& params) { #if defined(USE_NPU) - if (env_flag_enabled("XLLM_USE_ACLNN_LAYER_NORM_FWD", true)) { - return npu::layer_norm_fwd_aclnn(params.x, - params.weight, - params.bias, - params.eps, - params.z, - params.group_size, - params.norm_before_gate, - params.is_rms_norm); - } - return npu::layer_norm_fwd(params.x, - params.weight, - params.bias, - params.eps, - params.z, - params.group_size, - params.norm_before_gate, - params.is_rms_norm); + return npu::layer_norm_fwd_aclnn(params.x, + params.weight, + params.bias, + params.eps, + params.z, + params.group_size, + params.norm_before_gate, + params.is_rms_norm); #elif defined(USE_MLU) return mlu::gated_layer_norm(params.x, params.weight,