diff --git a/third_party/xllm_ops b/third_party/xllm_ops index bbbdfe3818..6aeb3fc06b 160000 --- a/third_party/xllm_ops +++ b/third_party/xllm_ops @@ -1 +1 @@ -Subproject commit bbbdfe3818149b6d4d725e1fbf44149c4c01a1a6 +Subproject commit 6aeb3fc06bf824d7dd0808c9061e91a87cc284da diff --git a/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt b/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt index 09bd02d541..1a6f071618 100644 --- a/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt +++ b/xllm/core/kernels/npu/xllm_ops/CMakeLists.txt @@ -34,6 +34,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 71f21f1348..fabbbc1ac1 100644 --- a/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h +++ b/xllm/core/kernels/npu/xllm_ops/xllm_ops_api.h @@ -374,4 +374,14 @@ std::pair npu_mega_chunk_gdn( bool output_final_state = false, const std::optional& cu_seqlens = std::nullopt, 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 fcb6fac575..56d3d9b556 100644 --- a/xllm/core/kernels/ops_api.cpp +++ b/xllm/core/kernels/ops_api.cpp @@ -35,6 +35,8 @@ limitations under the License. #include "dcu/dcu_ops_api.h" #endif +#include +#include #include #include "common/macros.h" @@ -47,6 +49,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) @@ -1595,6 +1615,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,