feat: make ATB context and operation cache device-aware for single-process mode.

custom_functions_npu/operation_create.h

@@ -60,7 +60,16 @@ atb::Operation* OpParamCache<ParamType>::get_operation(
   if (is_capturing) {
     return create_atb_operation(param, name);
   } else {
+    // Fold the current device into the cache key. An ATB operation created by
+    // atb::CreateOperation binds runner/workspace state to the device current
+    // at creation time, so a single-process multi-device run must not share one
+    // cached operation across devices (doing so makes Setup fail on the other
+    // device). Single-device runs keep one entry per param as before.
+    int32_t device_id = 0;
+    auto device_status = aclrtGetDevice(&device_id);
+    CHECK_EQ(device_status, ACL_ERROR_NONE) << "aclrtGetDevice failed!";
     uint64_t hashValue = compute_hash(param);
+    hashValue = hashValue * 131 + static_cast<uint64_t>(device_id) + 1;
     {
       std::lock_guard<std::mutex> lock(mutex_);
       auto op_cache = op_map_.find(hashValue);

custom_functions_npu/utils.cpp

@@ -10,24 +10,41 @@ ContextManager& ContextManager::get_instance() {
   return instance;
 }
 
-ContextManager::ContextManager() : atb_context_(nullptr) {}
+ContextManager::ContextManager() {}
 
 ContextManager::~ContextManager() {
-  if (atb_context_) {
-    auto status = atb::DestroyContext(atb_context_);
-    CHECK_EQ(status, 0) << "Destroy context failed!";
-    atb_context_ = nullptr;
+  for (auto& entry : device_contexts_) {
+    if (entry.second != nullptr) {
+      auto status = atb::DestroyContext(entry.second);
+      CHECK_EQ(status, 0) << "Destroy context failed!";
+    }
   }
+  device_contexts_.clear();
 }
 
 atb::Context* ContextManager::get_context(aclrtStream stream) {
-  std::call_once(create_flag_, [this]() {
-    auto status = atb::CreateContext(&atb_context_);
-    CHECK_EQ(status, 0) << "Create context failed!";
-  });
+  // ATB contexts are bound to the device current at creation time, so keep one
+  // per device. The calling thread runs on its worker's device (set via
+  // aclrtSetDevice), so query that device and create/reuse its context.
+  int32_t device_id = 0;
+  auto device_status = aclrtGetDevice(&device_id);
+  CHECK_EQ(device_status, ACL_ERROR_NONE) << "aclrtGetDevice failed!";
 
-  atb_context_->SetExecuteStream(stream);
-  return atb_context_;
+  atb::Context* context = nullptr;
+  {
+    std::lock_guard<std::mutex> lock(mutex_);
+    auto it = device_contexts_.find(device_id);
+    if (it == device_contexts_.end()) {
+      auto status = atb::CreateContext(&context);
+      CHECK_EQ(status, 0) << "Create context failed!";
+      device_contexts_[device_id] = context;
+    } else {
+      context = it->second;
+    }
+  }
+
+  context->SetExecuteStream(stream);
+  return context;
 }
 
 atb::Context* get_context(aclrtStream stream) {

custom_functions_npu/utils.h

@@ -5,11 +5,20 @@
 #include <glog/logging.h>
 #include <torch_npu/csrc/core/npu/NPUFormat.h>
 
+#include <mutex>
+#include <unordered_map>
+
 #include "atb/atb_infer.h"
 
 namespace atb {
 namespace utils {
 
+// Holds one ATB context per device. An ATB context is bound to the device that
+// is current when it is created, so a single shared context cannot be reused
+// across devices: in single-process multi-device mode each worker runs on its
+// own device, and reusing another device's context makes operation Setup fail.
+// Keying the context by device id keeps single-device behaviour unchanged (one
+// entry) while making multi-device single-process correct.
 class ContextManager {
  public:
   static ContextManager& get_instance();
@@ -21,8 +30,8 @@ class ContextManager {
 
  private:
   ContextManager();
-  std::once_flag create_flag_;
-  atb::Context* atb_context_;
+  std::mutex mutex_;
+  std::unordered_map<int32_t, atb::Context*> device_contexts_;
 };
 
 atb::Context* get_context(aclrtStream stream);

triton_npu/kernel_registry.cpp

@@ -21,7 +21,9 @@
 
 #include <filesystem>
 #include <fstream>
+#include <mutex>
 #include <nlohmann/json.hpp>
+#include <shared_mutex>
 
 namespace xllm::kernel::npu {
 
@@ -111,6 +113,10 @@ bool KernelRegistry::parse_json_config(const std::string& json_path,
 
 bool KernelRegistry::register_kernel(const std::string& kernel_name,
                                      const std::string& binary_path) {
+  // Registration mutates the shared map and performs process-global ACL
+  // registration; serialize it so concurrent per-device workers register each
+  // kernel exactly once and then share it.
+  std::unique_lock<std::shared_mutex> lock(mutex_);
   if (kernel_infos_.find(kernel_name) != kernel_infos_.end()) {
     LOG(INFO) << "Kernel '" << kernel_name << "' is already registered";
     return true;
@@ -188,6 +194,7 @@ bool KernelRegistry::register_kernel(const std::string& kernel_name,
 
 KernelStubHandle KernelRegistry::get_kernel_stub(
     const std::string& kernel_name) const {
+  std::shared_lock<std::shared_mutex> lock(mutex_);
   auto it = kernel_infos_.find(kernel_name);
   if (it == kernel_infos_.end()) {
     return nullptr;
@@ -197,13 +204,15 @@ KernelStubHandle KernelRegistry::get_kernel_stub(
 
 bool KernelRegistry::is_kernel_registered(
     const std::string& kernel_name) const {
+  std::shared_lock<std::shared_mutex> lock(mutex_);
   return kernel_infos_.find(kernel_name) != kernel_infos_.end();
 }
 
 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::shared_lock<std::shared_mutex> lock(mutex_);
   auto it = kernel_infos_.find(kernel_name);
   if (it != kernel_infos_.end()) {
     workspace_size = it->second.workspace_size;
@@ -296,6 +305,7 @@ bool KernelRegistry::register_binary(KernelInfo& info, uint32_t binary_size) {
 }
 
 void KernelRegistry::cleanup() {
+  std::unique_lock<std::shared_mutex> lock(mutex_);
   size_t count = kernel_infos_.size();
   for (auto& [name, info] : kernel_infos_) {
     if (info.buffer) {

triton_npu/kernel_registry.h

@@ -17,6 +17,7 @@
 #pragma once
 
 #include <cstdint>
+#include <shared_mutex>
 #include <string>
 #include <unordered_map>
 #include <utility>
@@ -113,6 +114,12 @@ class KernelRegistry {
   char* load_binary_file(const std::string& file_path, uint32_t& file_size);
   bool register_binary(KernelInfo& info, uint32_t binary_size);
 
+  // ACL kernel registration (rtDevBinaryRegister / rtFunctionRegister) is
+  // process-global, and kernel_infos_ is shared across all worker threads. In
+  // single-process multi-device mode one thread per device registers kernels
+  // concurrently, so all accesses must be serialized. Shared lock for reads,
+  // unique lock for register/cleanup.
+  mutable std::shared_mutex mutex_;
   std::unordered_map<std::string, KernelInfo> kernel_infos_;
 };