refactor: redesign wrapper for NPU fused_layernorm operator.

Author: yingxudeng

xllm/core/kernels/ops_api.cpp

@@ -283,6 +283,9 @@ void fused_layernorm(FusedLayerNormParams& params) {
                        params.store_output_before_norm,
                        params.store_output_after_norm,
                        params.dynamic_quant);
+#elif defined(USE_NPU)
+  params.output = npu::fused_layernorm(
+      params.input, params.weight, params.eps, params.mode);
 #elif defined(USE_CUDA)
   if (params.residual.has_value()) {
     cuda::fused_add_rms_norm(
@@ -306,15 +309,6 @@ void fused_layernorm(FusedLayerNormParams& params) {
 #endif
 }
 
-torch::Tensor fused_layernorm_tensor(FusedLayerNormParams& params) {
-#if defined(USE_NPU)
-  return npu::fused_layernorm(
-      params.input, params.weight, params.eps, params.mode);
-#else
-  LOG(FATAL) << "fused_layernorm not implemented";
-#endif
-}
-
 torch::Tensor matmul(MatmulParams& params) {
 #if defined(USE_MLU)
   return mlu::matmul(

xllm/core/kernels/ops_api.h

@@ -36,8 +36,6 @@ void batch_decode(AttentionParams& params);
 
 void fused_layernorm(FusedLayerNormParams& params);
 
-torch::Tensor fused_layernorm_tensor(FusedLayerNormParams& params);
-
 torch::Tensor matmul(MatmulParams& params);
 
 torch::Tensor group_gemm(GroupGemmParams& params);

xllm/core/layers/common/dense_mlp.cpp

@@ -18,6 +18,7 @@ limitations under the License.
 #include <glog/logging.h>
 
 #include "kernels/ops_api.h"
+#include "platform/device.h"
 
 namespace xllm {
 namespace layer {

xllm/core/layers/common/rms_norm.cpp

@@ -18,6 +18,7 @@ limitations under the License.
 #include <glog/logging.h>
 
 #include "kernels/ops_api.h"
+#include "platform/device.h"
 
 namespace xllm {
 namespace layer {
@@ -40,7 +41,10 @@ RMSNormImpl::RMSNormImpl(const ModelContext& context)
                   context.get_tensor_options()) {}
 
 torch::Tensor RMSNormImpl::forward(torch::Tensor& input) {
-  auto output = torch::empty_like(input);
+  torch::Tensor output;
+  if (Device::type_str() != "npu") {
+    output = torch::empty_like(input);
+  }
   return forward_output(input, output);
 }