[TransferEngine] Port URMA NUMA affinity mode

mooncake-store/include/client_service.h

@@ -662,6 +662,15 @@ class Client {
     // Return sorted NUMA node IDs that have at least one RDMA NIC.
     [[nodiscard]] std::vector<int> GetNicNumaNodes() const;
 
+    // Return the total number of NUMA nodes (counted from the local topology,
+    // which has one cpu:N entry per node regardless of NIC presence).
+    [[nodiscard]] int GetNumaNodeCount() const;
+
+    // Pre-establish transfer-engine connections to the currently registered
+    // segments using a small buffer from the client allocator.
+    [[nodiscard]] tl::expected<void, ErrorCode> warmup(
+        const std::shared_ptr<ClientBufferAllocator>& allocator);
+
     tl::expected<Replica::Descriptor, ErrorCode> GetPreferredReplica(
         const std::vector<Replica::Descriptor>& replica_list);
 

mooncake-store/src/client_service.cpp

@@ -3432,6 +3432,131 @@ std::vector<int> Client::GetNicNumaNodes() const {
     return {nodes.begin(), nodes.end()};
 }
 
+tl::expected<void, ErrorCode> Client::warmup(
+    const std::shared_ptr<ClientBufferAllocator>& allocator) {
+    if (!transfer_engine_) {
+        return tl::unexpected(ErrorCode::INTERNAL_ERROR);
+    }
+    if (!allocator) {
+        LOG(WARNING) << "warmup: no buffer allocator provided, skipping";
+        return {};
+    }
+
+    // Allocate a buffer from the client's buffer allocator.  The allocator's
+    // underlying memory is already registered with the transfer engine during
+    // setup_internal, so it can be used directly as the source (WRITE) and
+    // destination (READ) for warmup transfers.  BufferHandle is RAII: it
+    // deallocates the buffer back to the allocator on destruction, so no
+    // explicit register/unregister is needed here.
+    constexpr size_t kWarmupBufSize = 4096;
+    auto buf_handle = allocator->allocate(kWarmupBufSize);
+    if (!buf_handle) {
+        LOG(WARNING) << "warmup: failed to allocate warmup buffer";
+        return tl::unexpected(ErrorCode::INTERNAL_ERROR);
+    }
+    std::memset(buf_handle->ptr(), 0, buf_handle->size());
+
+    // Fetch all segment names registered with the master.
+    auto segments_result = master_client_.GetAllSegments();
+    if (!segments_result) {
+        LOG(WARNING) << "warmup: GetAllSegments failed: "
+                     << toString(segments_result.error());
+        return tl::unexpected(segments_result.error());
+    }
+
+    const auto& segments = segments_result.value();
+    LOG(INFO) << "warmup: pre-establishing connections to "
+              << segments.size() << " segment(s) for protocol '"
+              << protocol_ << "'";
+
+    // Helper lambda: submit a single transfer request and poll to completion.
+    auto submit_and_wait = [&](Transport::TransferRequest::OpCode op,
+                              SegmentID target_id,
+                              const std::string& segment_name) -> bool {
+        auto batch_id = transfer_engine_->allocateBatchID(1);
+        if (batch_id == INVALID_BATCH_ID) {
+            LOG(WARNING) << "warmup: allocateBatchID failed for '"
+                         << segment_name << "'";
+            return false;
+        }
+        auto target_meta_data = transfer_engine_->getMetadata();
+        auto target_segment = target_meta_data->getSegmentDescByID(target_id);
+        Transport::TransferRequest request;
+        request.opcode = op;
+        request.source = buf_handle->ptr();  // registered buffer carries r/w data
+        request.target_id = target_id;
+        request.target_offset = target_segment->buffers[0].addr;
+        request.length = kWarmupBufSize;
+
+        auto status = transfer_engine_->submitTransfer(batch_id, {request});
+        if (!status.ok()) {
+            LOG(WARNING) << "warmup: submitTransfer("
+                         << (op == Transport::TransferRequest::WRITE ? "W" : "R")
+                         << ") failed for '" << segment_name
+                         << "': " << status.message();
+            transfer_engine_->freeBatchID(batch_id);
+            return false;
+        }
+
+        // Poll for completion — the connection handshake occurs here.
+        Transport::TransferStatus ts;
+        for (int poll = 0; poll < 1000; ++poll) {
+            auto s = transfer_engine_->getTransferStatus(batch_id, 0, ts);
+            if (!s.ok()) break;
+            if (ts.s == Transport::COMPLETED ||
+                ts.s == Transport::FAILED ||
+                ts.s == Transport::INVALID)
+                break;
+            std::this_thread::sleep_for(std::chrono::microseconds(100));
+        }
+
+        transfer_engine_->freeBatchID(batch_id);
+        return ts.s == Transport::COMPLETED;
+    };
+
+    size_t success_count = 0;
+    for (const auto& segment_name : segments) {
+        // Open the segment to resolve its SegmentID (caches the descriptor).
+        auto target_id = transfer_engine_->openSegment(segment_name);
+        if (target_id == (SegmentHandle)-1) {
+            LOG(WARNING) << "warmup: cannot open segment '"
+                         << segment_name << "'";
+            continue;
+        }
+        // Skip the local segment — no point self-connecting, and a WRITE
+        // would corrupt our own buffer.
+        if (target_id == LOCAL_SEGMENT_ID) {
+            continue;
+        }
+
+        // Issue a 1-byte WRITE then a 1-byte READ using the registered
+        // buffer to exercise both directions of the connection.
+        bool ok = submit_and_wait(Transport::TransferRequest::WRITE,
+                                  target_id, segment_name);
+        ok = submit_and_wait(Transport::TransferRequest::READ,
+                             target_id, segment_name) && ok;
+        if (ok) ++success_count;
+    }
+
+    LOG(INFO) << "warmup: completed, " << success_count << "/"
+              << segments.size() << " segments processed";
+    return {};
+}
+
+int Client::GetNumaNodeCount() const {
+    if (!transfer_engine_) return 0;
+    auto topo = transfer_engine_->getLocalTopology();
+    if (!topo) return 0;
+    // discoverCpuTopology emits one "cpu:N" entry per NUMA node (regardless of
+    // whether it hosts a NIC), so counting them gives the NUMA node count.
+    int count = 0;
+    for (auto& [name, entry] : topo->getMatrix()) {
+        (void)entry;
+        if (name.rfind("cpu:", 0) == 0) ++count;
+    }
+    return count;
+}
+
 tl::expected<void, ErrorCode> Client::MountSegment(
     const void* buffer, size_t size, const std::string& protocol,
     const std::string& location) {

mooncake-store/src/real_client.cpp

@@ -40,6 +40,9 @@
 #ifdef USE_NOF
 #include "spdk/spdk_wrapper.h"
 #endif
+#ifdef USE_UB
+#include "ub_allocator.h"
+#endif
 #ifdef USE_ASCEND_DIRECT
 #include "acl/acl_rt.h"
 #include "transport/ascend_transport/ascend_direct_transport/context_manager.h"
@@ -951,7 +954,8 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
 #endif
 
         // For RDMA, auto-discover NUMA nodes with NICs and distribute
-        // global_segment across them for full NIC utilization.
+        // global_segment across them for full NIC utilization. RDMA keeps the
+        // legacy single-segment-with-multi-region ("segments:...") behavior.
         std::vector<int> seg_numa_nodes;
         if (protocol == "rdma") {
             seg_numa_nodes = client_->GetNicNumaNodes();
@@ -968,6 +972,32 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
             }
         }
 
+        // For UB, mount ONE segment per NUMA node (each bound to its node,
+        // location="cpu:N"). Memory is spread across ALL NUMA nodes by count,
+        // not just NIC-bearing ones: this works even with a single bonded
+        // device (e.g. bonding_dev_0, NUMA=-1) where NIC-NUMA discovery would
+        // be empty, and it relies only on the (decimal) NUMA count, so it is
+        // unaffected by the hex "numa" attribute. Each segment then drives
+        // selectDevice (cpu:N -> local NIC, else any) and chip affinity
+        // (cpu:N -> chip via numaNodeToChipId). Automatic whenever UB has more
+        // than one NUMA node; independent of both MC_UB_NUMA_AFFINITY_ENABLE
+        // and MC_URMA_BONDING_MULTIPATH_ENABLE.
+#ifdef USE_UB
+        std::vector<int> ub_numa_nodes;
+        if (protocol == "ub") {
+            int numa_count = client_->GetNumaNodeCount();
+            if (numa_count > 1) {
+                std::string nodes_str;
+                for (int i = 0; i < numa_count; ++i) {
+                    ub_numa_nodes.push_back(i);
+                    if (i) nodes_str += ",";
+                    nodes_str += std::to_string(i);
+                }
+                MC_LOG(INFO) << "UB per-NUMA mode: NUMA node count=" << numa_count
+                          << ", nodes=[" << nodes_str << "]";
+            }
+        }
+#endif  // USE_UB
         const bool parallel_hugetlb_population =
             protocol == "rdma" && should_use_hugepage;
 
@@ -981,6 +1011,54 @@ tl::expected<void, ErrorCode> RealClient::setup_internal(
             }
             global_segment_size -= segment_size;
 
+            // UB NUMA affinity: split this chunk into one segment per NIC-NUMA
+            // node, each physically bound to its node and registered with
+            // location "cpu:N" (so selectDevice picks the NUMA-local NIC).
+#ifdef USE_UB
+            if (!ub_numa_nodes.empty()) {
+                size_t page_sz = should_use_hugepage
+                                     ? get_hugepage_size_from_env()
+                                     : static_cast<size_t>(getpagesize());
+                size_t n = ub_numa_nodes.size();
+                size_t per_node_size = align_up(segment_size / n, page_sz);
+                if (per_node_size == 0) {
+                    MC_LOG(ERROR) << "UB per-NUMA: per_node_size is 0, segment "
+                                     "too small for " << n << " NUMA nodes";
+                    return tl::unexpected(ErrorCode::INVALID_PARAMS);
+                }
+                for (int node : ub_numa_nodes) {
+                    // Use UB's own allocator bound to this node: numa_alloc_onnode
+                    // via libnuma, registered in the store-memory table, so URMA
+                    // can register it. (A raw mmap+mbind buffer cannot be
+                    // registered by urma_register_seg -- it fails with error
+                    // 2048 because the VMA has no backing pages at reg time.)
+                    void *ptr = mooncake::ub_allocate_memory_onnode(
+                        /*alignment=*/page_sz, per_node_size, node);
+                    if (!ptr) {
+                        MC_LOG(ERROR) << "UB per-NUMA: failed to allocate "
+                                         "segment for node " << node;
+                        return tl::unexpected(ErrorCode::INVALID_PARAMS);
+                    }
+                    // numa_alloc-backed => free via ub_free_memory/numa_free,
+                    // NOT munmap. Track with UbSegmentDeleter accordingly.
+                    ub_segment_ptrs_.emplace_back(
+                        ptr, UbSegmentDeleter{per_node_size});
+
+                    std::string loc = genCpuNodeName(node);  // "cpu:<node>"
+                    MC_LOG(INFO) << "Mounting UB per-NUMA segment: node=" << node
+                              << ", size=" << per_node_size << ", loc=" << loc;
+                    auto mr = client_->MountSegment(ptr, per_node_size, protocol,
+                                                    loc);
+                    if (!mr.has_value()) {
+                        MC_LOG(ERROR) << "Failed to mount UB per-NUMA segment: "
+                                   << toString(mr.error());
+                        return tl::unexpected(mr.error());
+                    }
+                }
+                continue;  // this chunk fully mounted across NUMA nodes
+            }
+#endif  // USE_UB
+
             size_t mapped_size = segment_size;
             void *ptr = nullptr;
             std::string seg_location = kWildcardLocation;

mooncake-transfer-engine/include/config.h

@@ -167,6 +167,14 @@ struct GlobalConfig {
     uint64_t max_seg_size = 0x10000000000;
     size_t max_jfc_e = 4096;  // urma is temporarily using this default value.
     size_t num_jetty_per_ep = 1;
+    // Enable URMA bonding multipath mode. Default is off; override via
+    // MC_URMA_BONDING_MULTIPATH_ENABLE.
+    bool urma_bonding_multipath = false;
+    // Enable UB NUMA affinity: store splits the global segment into one
+    // segment per NIC-NUMA node, and transfers pin src/dst chip by NUMA.
+    // Independent from urma_bonding_multipath; default off; override via
+    // MC_UB_NUMA_AFFINITY_ENABLE.
+    bool ub_numa_affinity = false;
 };
 
 struct RpcCommunicatorConfig {

mooncake-transfer-engine/include/memory_location.h

@@ -31,6 +31,18 @@ struct MemoryLocationEntry {
     std::string location;
 };
 
+const static uint8_t INVALID_CHIP_ID = 0xFF;
+
+// "cpu:3" -> 3; "*" 或 非cpu串返回-1
+int parseCpuNumaNode(const std::string& location);
+
+// NUMA 节点 -> chip id：优先按 sysfs 真实拓扑(physical_package_id)映射；
+// sysfs 不可读时回退「前半 chip1 / 后半 chip2」启发式；失败返回 INVALID_CHIP_ID。
+uint8_t numaNodeToChipId(int numa_node, size_t numa_count = 0);
+
+// NUMA 节点 -> location 字符串："cpu:N"（node>=0）或 "*"（node<0）。
+std::string genCpuNodeName(int node);
+
 // If only_first_page is true, only the location of the first page will be
 // returned. Scan all pages may take a long time, so set only_first_page if only
 // the location of the first page is needed.

mooncake-transfer-engine/include/transfer_metadata.h

@@ -75,6 +75,11 @@ class TransferMetadata {
         uint64_t offset;                    // for cxl
         std::vector<std::string> tseg;      // for ub/urma
         std::vector<uint32_t> l_seg_index;  // for ub/urma
+        // for ub: NUMA->chip id of this buffer, computed once by the owner at
+        // registration and published. -1 = not provided (consumer falls back
+        // to resolving from `name`). Only meaningful for single-NUMA ("cpu:N")
+        // buffers; multi-NUMA "segments:..." buffers leave it -1.
+        int chip_id = -1;
 
         bool operator==(const BufferDesc &other) const = default;
     };

mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h

@@ -138,7 +138,9 @@ class UbContext {
           max_endpoints_(max_endpoints),
           worker_pool_(nullptr),
           active_(true),
-          show_work_request_flushed_error_(false) {}
+          show_work_request_flushed_error_(false),
+          multipath_(globalConfig().urma_bonding_multipath),
+          numa_affinity_(globalConfig().ub_numa_affinity) {}
 
     virtual ~UbContext() = default;
 
@@ -214,6 +216,10 @@ class UbContext {
 
     void set_active(bool flag) { active_ = flag; }
 
+    bool multipath() const { return multipath_; }
+
+    bool numa_affinity() const { return numa_affinity_; }
+
     // EndPoint Management
     std::shared_ptr<UbEndPoint> endpoint() {
         return endpoint("LOCAL_SEGMENT_ID");
@@ -369,6 +375,10 @@ class UbContext {
     volatile bool active_;
 
     bool show_work_request_flushed_error_;
+
+    bool multipath_ = false;
+
+    bool numa_affinity_ = false;
 };
 }  // namespace mooncake
 

mooncake-transfer-engine/include/transport/transport.h

@@ -161,6 +161,8 @@ class Transport {
                 uint32_t max_retry_cnt;
                 void *r_seg;
                 void *l_seg;
+                uint8_t src_chip_id;
+                uint8_t dst_chip_id;
                 void *endpoint;
             } ub;
             struct {

mooncake-transfer-engine/include/ub_allocator.h

@@ -4,6 +4,12 @@ namespace mooncake {
 
 void* ub_allocate_memory(size_t alignment, size_t total_size);
 
+// Same as ub_allocate_memory but binds the allocation to a specific NUMA node
+// (numa_node < 0 falls back to node-local). Still allocated via libnuma and
+// registered in the store-memory range table, so URMA can register it.
+void* ub_allocate_memory_onnode(size_t alignment, size_t total_size,
+                                int numa_node);
+
 void ub_free_memory(void* ptr);
 
 bool ub_is_store_memory(void* addr, size_t length);

mooncake-transfer-engine/src/config.cpp

@@ -728,6 +728,30 @@ void loadGlobalConfig(GlobalConfig& config) {
         }
     }
 
+    const char* urma_bonding_multipath_enable =
+        std::getenv("MC_URMA_BONDING_MULTIPATH_ENABLE");
+    if (urma_bonding_multipath_enable && *urma_bonding_multipath_enable) {
+        std::string val(urma_bonding_multipath_enable);
+        if (val == "true" || val == "1" || val == "on") {
+            config.urma_bonding_multipath = true;
+            LOG(WARNING) << "MC_URMA_BONDING_MULTIPATH_ENABLE is " << val;
+        } else
+            LOG(WARNING)
+                << "Ignore value from environment variable "
+                   "MC_URMA_BONDING_MULTIPATH_ENABLE, it should be true|1|on";
+    }
+
+    const char* ub_numa_affinity_enable =
+        std::getenv("MC_UB_NUMA_AFFINITY_ENABLE");
+    if (ub_numa_affinity_enable && *ub_numa_affinity_enable) {
+        std::string val(ub_numa_affinity_enable);
+        if (val == "true" || val == "1" || val == "on") {
+            config.ub_numa_affinity = true;
+            LOG(WARNING) << "MC_UB_NUMA_AFFINITY_ENABLE is " << val;
+        } else
+            LOG(WARNING) << "Ignore value from environment variable "
+                            "MC_UB_NUMA_AFFINITY_ENABLE, it should be true|1|on";
+    }
     const char* mlx5_qp_lag_port_balance_env =
         std::getenv("MC_MLX5_QP_LAG_PORT_BALANCE");
     if (mlx5_qp_lag_port_balance_env && *mlx5_qp_lag_port_balance_env) {