From 21cab2b6af14f18f56b7453024319065c8fb9764 Mon Sep 17 00:00:00 2001
From: pjgao <gaopengju3@huawei.com>
Date: Tue, 21 Jul 2026 12:28:16 +0800
Subject: [PATCH] fix(npu): serialize kernel registry access

---
 triton_npu/kernel_registry.cpp | 15 +++++++++++++++
 1 file changed, 15 insertions(+)

diff --git a/triton_npu/kernel_registry.cpp b/triton_npu/kernel_registry.cpp
index 7382be5..60c2ce8 100644
--- a/triton_npu/kernel_registry.cpp
+++ b/triton_npu/kernel_registry.cpp
@@ -21,10 +21,20 @@
 
 #include <filesystem>
 #include <fstream>
+#include <mutex>
 #include <nlohmann/json.hpp>
 
 namespace xllm::kernel::npu {
 
+namespace {
+
+std::mutex& registry_mutex() {
+  static std::mutex mutex;
+  return mutex;
+}
+
+}  // namespace
+
 KernelRegistry& KernelRegistry::get_instance() {
   static KernelRegistry instance;
   return instance;
@@ -111,6 +121,7 @@ bool KernelRegistry::parse_json_config(const std::string& json_path,
 
 bool KernelRegistry::register_kernel(const std::string& kernel_name,
                                      const std::string& binary_path) {
+  std::lock_guard<std::mutex> lock(registry_mutex());
   if (kernel_infos_.find(kernel_name) != kernel_infos_.end()) {
     LOG(INFO) << "Kernel '" << kernel_name << "' is already registered";
     return true;
@@ -188,6 +199,7 @@ bool KernelRegistry::register_kernel(const std::string& kernel_name,
 
 KernelStubHandle KernelRegistry::get_kernel_stub(
     const std::string& kernel_name) const {
+  std::lock_guard<std::mutex> lock(registry_mutex());
   auto it = kernel_infos_.find(kernel_name);
   if (it == kernel_infos_.end()) {
     return nullptr;
@@ -197,6 +209,7 @@ KernelStubHandle KernelRegistry::get_kernel_stub(
 
 bool KernelRegistry::is_kernel_registered(
     const std::string& kernel_name) const {
+  std::lock_guard<std::mutex> lock(registry_mutex());
   return kernel_infos_.find(kernel_name) != kernel_infos_.end();
 }
 
@@ -204,6 +217,7 @@ bool KernelRegistry::get_kernel_workspace_config(const std::string& kernel_name,
                                                  int64_t& workspace_size,
                                                  int64_t& lock_init_value,
                                                  int64_t& lock_num) const {
+  std::lock_guard<std::mutex> lock(registry_mutex());
   auto it = kernel_infos_.find(kernel_name);
   if (it != kernel_infos_.end()) {
     workspace_size = it->second.workspace_size;
@@ -296,6 +310,7 @@ bool KernelRegistry::register_binary(KernelInfo& info, uint32_t binary_size) {
 }
 
 void KernelRegistry::cleanup() {
+  std::lock_guard<std::mutex> lock(registry_mutex());
   size_t count = kernel_infos_.size();
   for (auto& [name, info] : kernel_infos_) {
     if (info.buffer) {