refactor: redesign wrapper for NPU fused_layernorm operator.

Author: yingxudeng

xllm/core/kernels/ops_api.cpp

@@ -229,14 +229,8 @@ void fused_layernorm(FusedLayerNormParams& params) {
                        params.dynamic_quant);
 #elif defined(USE_CUDA)
   cuda::rmsnorm(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 fused_moe(FusedMoEParams& 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);
 }