refactor: move NPULayerSynchronizerImpl to platform dir. (jd-opensource#237)

Author: XuZhang99

xllm/core/common/CMakeLists.txt

@@ -8,7 +8,6 @@ cc_library(
     etcd_client.h
     global_flags.h
     instance_name.h
-    layer_synchronizer.h
     macros.h
     metrics.h
     $<$<BOOL:${USE_NPU}>:mspti_helper.h>
@@ -19,7 +18,6 @@ cc_library(
   SRCS
     etcd_client.cpp
     global_flags.cpp
-    layer_synchronizer.cpp
     metrics.cpp
     $<$<BOOL:${USE_NPU}>:mspti_helper.cpp>
     options.cpp

xllm/core/distributed_runtime/worker_service.cpp

@@ -388,7 +388,7 @@ void WorkerService::ExecuteModel(
             top_logprobs =
                 safe_to(sample_output.top_logprobs, torch::kCPU, true);
           }
-          auto ret = stream_->synchronize_stream();
+          auto ret = stream_->synchronize();
         }
       }
     } else {
@@ -460,7 +460,7 @@ void WorkerService::GetLastStepResult(
             // [num_seq, topk]
             const auto& top_logprobs =
                 safe_to(sample_output.top_logprobs, torch::kCPU, true);
-            auto ret = stream_->synchronize_stream();
+            auto ret = stream_->synchronize();
 
             forward_output_to_proto(next_tokens,
                                     logprobs,

xllm/core/framework/eplb/eplb_executor.cpp

@@ -94,7 +94,7 @@ void EplbExecutor::eplb_worker_loop() {
 
     c10::StreamGuard streamGuard = stream_->set_stream_guard();
     model_->prepare_expert_weight(task.layer_id, task.expert_ids);
-    auto ret = stream_->synchronize_stream();
+    auto ret = stream_->synchronize();
     auto prepare_end = std::chrono::high_resolution_clock::now();
     auto prepare_duration =
         std::chrono::duration_cast<std::chrono::milliseconds>(prepare_end -

xllm/core/framework/kv_cache/CMakeLists.txt

@@ -32,4 +32,5 @@ cc_library(
     $<$<BOOL:${USE_NPU}>:torch_npu>
     mooncake_store
     :xtensor
+    $<$<BOOL:${USE_NPU}>:platform_npu>
 )

xllm/core/framework/kv_cache/kv_cache_transfer.h

@@ -17,10 +17,10 @@ limitations under the License.
 
 #include <folly/futures/Future.h>
 
-#include "common/layer_synchronizer.h"
 #include "common/types.h"
 #include "core/framework/parallel_state.h"
 #include "kv_cache.h"
+#include "platform/npu/npu_layer_synchronizer.h"
 #include "util/threadpool.h"
 
 namespace xllm {

xllm/core/framework/model/CMakeLists.txt

@@ -18,6 +18,7 @@ if(USE_NPU)
   # Modify dependencies for npu
   list(APPEND BASE_DEPS torch_npu)
   list(APPEND BASE_DEPS :npu_layers)
+  list(APPEND BASE_DEPS :platform_npu)
 endif()
 
 # Define the library

xllm/core/framework/model/model_input_params.h

@@ -19,7 +19,7 @@ limitations under the License.
 #include <torch/torch.h>
 
 #if defined(USE_NPU)
-#include "common/layer_synchronizer.h"
+#include "platform/npu/npu_layer_synchronizer.h"
 #endif
 #include "framework/request/mm_data.h"
 #include "npu_dp_ep_padding.h"

xllm/core/platform/CMakeLists.txt

@@ -13,4 +13,8 @@ cc_library(
     torch
     $<$<BOOL:${USE_NPU}>:torch_npu>
     $<$<BOOL:${USE_NPU}>:ascendcl>
-)
+)
+
+if(USE_NPU)
+  add_subdirectory(npu)
+endif()

xllm/core/platform/npu/CMakeLists.txt

@@ -0,0 +1,15 @@
+include(cc_library)
+
+cc_library(
+  NAME 
+    platform_npu
+  HDRS
+    npu_layer_synchronizer.h
+  SRCS
+    npu_layer_synchronizer.cpp
+  DEPS
+    torch_npu
+    glog::glog
+    torch
+    ascendcl
+)

xllm/core/common/layer_synchronizer.cpp renamed to xllm/core/platform/npu/npu_layer_synchronizer.cpp

@@ -13,13 +13,12 @@ See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
 
-#include "common/layer_synchronizer.h"
+#include "npu_layer_synchronizer.h"
 
 #include <glog/logging.h>
 
 namespace xllm {
 
-#if defined(USE_NPU)
 NPULayerSynchronizerImpl::NPULayerSynchronizerImpl(const int64_t num_layers)
     : events_(num_layers, nullptr), event_record_flags_(num_layers) {
   uint32_t flags = ACL_EVENT_SYNC;
@@ -53,6 +52,5 @@ bool NPULayerSynchronizerImpl::synchronize_layer(const int64_t layer_index) {
   }
   return true;
 }
-#endif
 
 }  // namespace xllm