refactor: redesign wrapper for NPU fused_layernorm operator.

Author: yingxudeng

xllm/core/kernels/ops_api.cpp

@@ -216,6 +216,10 @@ void batch_decode(AttentionParams& params) {
 
 void fused_layernorm(FusedLayerNormParams& params) {
 #if defined(USE_MLU)
+  params.output = torch::empty(
+      {params.input.sizes()[0], params.intermediate_size / params.world_size},
+      params.input.options());
+
   mlu::fused_layernorm(params.input,
                        params.output,
                        params.residual,
@@ -238,16 +242,14 @@ void fused_layernorm(FusedLayerNormParams& params) {
     params.output = params.input;
     params.residual_out = params.residual;
   } else {
+    params.output = torch::empty(
+        {params.input.sizes()[0], params.intermediate_size / params.world_size},
+        params.input.options());
+
     cuda::rms_norm(params.output, params.input, params.weight, params.eps);
   }
-#else
-  LOG(FATAL) << "fused_layernorm not implemented";
-#endif
-}
-
-torch::Tensor fused_layernorm_tensor(FusedLayerNormParams& params) {
-#if defined(USE_NPU)
-  return npu::fused_layernorm(
+#elif defined(USE_NPU)
+  params.output = npu::fused_layernorm(
       params.input, params.weight, params.eps, params.mode);
 #else
   LOG(FATAL) << "fused_layernorm not implemented";

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/fuse_norm.cpp

@@ -35,7 +35,11 @@ FusedRMSNormImpl::FusedRMSNormImpl(int64_t dim,
 }
 
 torch::Tensor FusedRMSNormImpl::forward(torch::Tensor& input) {
+#if defined(USE_NPU)
+  torch::Tensor output;
+#else
   auto output = torch::empty_like(input);
+#endif
   return forward_output(input, output);
 }