refactor: extend NPU kernel parameter set and adjust conditional compilation structure.

Author: yingxudeng

xllm/core/common/global_flags.cpp

@@ -390,4 +390,4 @@ DEFINE_string(reasoning_parser,
 // --- qwen3 reranker config ---
 DEFINE_bool(enable_qwen3_reranker, false, "Whether to enable qwen3 reranker.");
 
-DEFINE_bool(enable_native_npu, true, "Whether to enable native NPU support.");
+DEFINE_bool(enable_npu_torch, true, "Whether to enable native NPU support.");

xllm/core/common/global_flags.h

@@ -203,4 +203,4 @@ DECLARE_string(reasoning_parser);
 
 DECLARE_bool(enable_shm);
 
-DECLARE_bool(enable_native_npu);
+DECLARE_bool(enable_npu_torch);

xllm/core/distributed_runtime/worker_server.cpp

@@ -103,7 +103,7 @@ void WorkerServer::create_server(
 #elif defined(USE_NPU)
   // TODO: Refactor to use model_type or other appropriate enumeration for
   // condition checking
-  if (FLAGS_enable_native_npu) {
+  if (FLAGS_enable_npu_torch) {
     comm.create_process_groups(master_node_addr, device);
   }
 #endif

xllm/core/framework/parallel_state/collective_communicator.cpp

@@ -18,8 +18,6 @@ limitations under the License.
 #include "mapping_npu.h"
 
 #if defined(USE_NPU)
-#include <torch_npu/csrc/distributed/ProcessGroupHCCL.hpp>
-
 #include "npu_process_group.h"
 #include "xllm_kernels/core/include/atb_speed/base/external_comm_manager.h"
 #include "xllm_kernels/core/include/atb_speed/utils/singleton.h"

xllm/core/framework/parallel_state/npu_process_group.h

@@ -18,10 +18,6 @@ limitations under the License.
 #include "hccl/hccl.h"
 #include "process_group.h"
 
-namespace c10d_npu {
-class ProcessGroupHCCL;
-}
-
 namespace xllm {
 
 class ProcessGroupHCCL : public ProcessGroup {

xllm/core/framework/parallel_state/process_group.h

@@ -66,6 +66,9 @@ class ProcessGroup {
 
  protected:
 #if defined(USE_NPU)
+  // Using ProcessGroupHCCL for NPU devices
+  // Note: torch_npu uses an older torch version where c10d::Backend lacks
+  // shutdown() method
   std::unique_ptr<c10d_npu::ProcessGroupHCCL> pg_{nullptr};
 #else
   std::unique_ptr<c10d::Backend> pg_{nullptr};

xllm/core/kernels/npu/active.cpp

@@ -20,7 +20,11 @@ limitations under the License.
 
 namespace xllm::kernel::npu {
 
-torch::Tensor active(const torch::Tensor& input) {
+torch::Tensor active(const torch::Tensor& input, const std::string& act_mode) {
+  if (act_mode != "silu" && act_mode != "swiglu") {
+    throw std::runtime_error(
+        "Only swiglu activation is supported in NPU active");
+  }
   return at_npu::native::custom_ops::npu_swiglu(input);
 }
 }  // namespace xllm::kernel::npu

xllm/core/kernels/npu/attention.cpp

@@ -31,31 +31,34 @@ void batch_prefill(const torch::Tensor& query,
                    const torch::Tensor& mask,
                    const torch::Tensor& seq_len,
                    float scale,
-                   int num_heads,
-                   int num_kv_heads,
                    torch::Tensor& output) {
+  auto num_heads = query.size(-2);
+  auto num_kv_heads = key.size(-2);
   atb::_npu_flash_attention(
       query, key, value, mask, seq_len, scale, num_heads, num_kv_heads, output);
 }
 
 void batch_decode(const torch::Tensor& query,
                   const torch::Tensor& k_cache,
                   const torch::Tensor& v_cache,
-                  int num_kv_heads,
-                  int num_heads,
                   float scale,
                   const torch::Tensor& block_table,
                   const torch::Tensor& seq_lens,
                   torch::Tensor& output) {
-  atb::_npu_paged_attention(query,
+  auto head_size = query.size(-1);
+  auto num_heads = query.size(-2);
+  auto num_kv_heads = k_cache.size(-2);
+  auto q = query.view({-1, num_heads, head_size});
+  auto o = output.view({-1, num_heads, head_size});
+  atb::_npu_paged_attention(q,
                             k_cache,
                             v_cache,
                             num_kv_heads,
                             num_heads,
                             scale,
                             block_table,
                             seq_lens,
-                            output);
+                            o);
 }
 
 }  // namespace xllm::kernel::npu

xllm/core/kernels/npu/fused_layernorm.cpp

@@ -21,7 +21,12 @@ namespace xllm::kernel::npu {
 
 torch::Tensor fused_layernorm(const torch::Tensor& input,
                               const torch::Tensor& weight,
-                              double eps) {
+                              double eps,
+                              const std::string& mode) {
+  if (mode != "rmsnorm") {
+    throw std::runtime_error(
+        "Only rmsnorm mode is supported in NPU fused_layernorm");
+  }
   std::tuple<at::Tensor, at::Tensor> result =
       at_npu::native::custom_ops::npu_rms_norm(input, weight, eps);
   auto normalized_input = std::get<0>(result);

xllm/core/kernels/npu/npu_ops_api.h

@@ -34,15 +34,11 @@ void batch_prefill(const torch::Tensor& query,
                    const torch::Tensor& mask,
                    const torch::Tensor& seq_len,
                    float scale,
-                   int num_heads,
-                   int num_kv_heads,
                    torch::Tensor& output);
 
 void batch_decode(const torch::Tensor& query,
                   const torch::Tensor& k_cache,
                   const torch::Tensor& v_cache,
-                  int num_kv_heads,
-                  int num_heads,
                   float scale,
                   const torch::Tensor& block_table,
                   const torch::Tensor& seq_lens,
@@ -52,11 +48,12 @@ torch::Tensor matmul(const torch::Tensor& a,
                      const torch::Tensor& b,
                      const std::optional<torch::Tensor>& bias);
 
-torch::Tensor active(const torch::Tensor& input);
+torch::Tensor active(const torch::Tensor& input, const std::string& act_mode);
 
 torch::Tensor fused_layernorm(const torch::Tensor& input,
                               const torch::Tensor& weight,
-                              double eps);
+                              double eps,
+                              const std::string& mode);
 
 void apply_rotary(torch::Tensor& q,
                   torch::Tensor& k,