From e845567ed8a1b39abd3a7b612cab469555d1eb3c Mon Sep 17 00:00:00 2001 From: victoryang00 Date: Fri, 7 Aug 2026 18:41:07 -0700 Subject: [PATCH 1/4] Support 16-node distributed Soft-RoCE routing --- include/distributed_server.h | 19 ++- include/rdma_communication.h | 12 +- src/distributed_server.cpp | 118 +++++++++++++++--- src/main_server.cc | 11 +- src/rdma_communication.cpp | 231 ++++++++++++++++++++++++----------- 5 files changed, 299 insertions(+), 92 deletions(-) diff --git a/include/distributed_server.h b/include/distributed_server.h index 0245286..8b159cf 100644 --- a/include/distributed_server.h +++ b/include/distributed_server.h @@ -402,6 +402,8 @@ struct RDMACalibrationResult { /* Per-node RDMA connection state */ struct RDMANodeConnection { std::unique_ptr client; + std::string endpoint_addr; + uint16_t endpoint_port; // Outgoing RDMA connection uint64_t remote_addr; // Remote base address @@ -412,7 +414,7 @@ struct RDMANodeConnection { RDMACalibrationResult calibration; // Per-node calibration data - RDMANodeConnection() : remote_addr(0), remote_buffer_size(0), connected(false) {} + RDMANodeConnection() : endpoint_port(0), remote_addr(0), remote_buffer_size(0), connected(false) {} }; /* Message handler callback type */ @@ -646,6 +648,9 @@ class DistributedRDMATransport { // RDMA server for incoming connections std::unique_ptr server_; std::thread accept_thread_; + std::vector> incoming_connections_; + std::vector incoming_threads_; + std::mutex incoming_mutex_; std::atomic running_; // Calibration results per node @@ -660,8 +665,18 @@ class DistributedRDMATransport { bool initialize(); void shutdown(); + // Route incoming two-sided RDMA requests to the distributed server's + // local-memory implementation. Without this, RDMAServer's default + // handler only returns a synthetic success response. + void set_message_handler(RDMAConnection::MessageHandler handler) { + if (server_) { + server_->set_message_handler(handler); + } + } + // Connection management - bool connect_to_node(uint32_t node_id, const std::string &addr, uint16_t port); + bool connect_to_node(uint32_t node_id, const std::string &addr, uint16_t port, uint64_t remote_addr, + size_t remote_buffer_size); void disconnect_node(uint32_t node_id); bool is_connected(uint32_t node_id) const; std::vector get_connected_nodes() const; diff --git a/include/rdma_communication.h b/include/rdma_communication.h index badc9b3..ef22441 100644 --- a/include/rdma_communication.h +++ b/include/rdma_communication.h @@ -4,6 +4,7 @@ #include #include #include +#include #include #include @@ -83,9 +84,13 @@ class RDMAConnection { RDMAConnection(); virtual ~RDMAConnection(); +#ifdef HAS_RDMA + int accept_cm_id(struct rdma_cm_id *id); +#endif + void mark_connected(); void set_message_handler(MessageHandler handler) { message_handler_ = handler; } int send_message(const RDMAMessage &msg); - int receive_message(RDMAMessage &msg); + int receive_message(RDMAMessage &msg, int timeout_ms = -1); bool is_connected() const { return connected_.load(); } void disconnect(); }; @@ -96,6 +101,7 @@ class RDMAServer : public RDMAConnection { uint16_t port_; #ifdef HAS_RDMA struct rdma_cm_id *listen_id_; + std::map> pending_connections_; #endif public: @@ -103,8 +109,8 @@ class RDMAServer : public RDMAConnection { ~RDMAServer(); int start(); - int accept_connection(); - void handle_client(); + std::shared_ptr accept_connection(); + void handle_client(const std::shared_ptr &client); void stop(); }; diff --git a/src/distributed_server.cpp b/src/distributed_server.cpp index 4861b1c..286241a 100644 --- a/src/distributed_server.cpp +++ b/src/distributed_server.cpp @@ -557,8 +557,11 @@ bool DistributedMemoryServer::initialize() { // Initialize message manager msg_manager_ = std::make_unique(shm_name_, node_id_); - bool is_first_node = (node_id_ == 0); - if (!msg_manager_->initialize(is_first_node)) { + // POSIX SHM is host-local. Only SHM transport can join node 0's control + // segment directly; physical TCP/RDMA nodes need an independent local + // control segment and exchange data over their selected transport. + bool create_local_control = (node_id_ == 0 || transport_mode_ != DistTransportMode::SHM); + if (!msg_manager_->initialize(create_local_control)) { SPDLOG_ERROR("Failed to initialize message manager"); return false; } @@ -1957,10 +1960,25 @@ void DistributedRDMATransport::shutdown() { server_->stop(); } + { + std::lock_guard incoming_lock(incoming_mutex_); + for (auto &client : incoming_connections_) { + if (client) + client->disconnect(); + } + } + if (accept_thread_.joinable()) { accept_thread_.join(); } + for (auto &thread : incoming_threads_) { + if (thread.joinable()) + thread.join(); + } + incoming_threads_.clear(); + incoming_connections_.clear(); + std::lock_guard lock(connections_mutex_); for (auto &[node_id, conn] : connections_) { if (conn.client) { @@ -1971,7 +1989,8 @@ void DistributedRDMATransport::shutdown() { connections_.clear(); } -bool DistributedRDMATransport::connect_to_node(uint32_t node_id, const std::string &addr, uint16_t port) { +bool DistributedRDMATransport::connect_to_node(uint32_t node_id, const std::string &addr, uint16_t port, + uint64_t remote_addr, size_t remote_buffer_size) { std::lock_guard lock(connections_mutex_); auto it = connections_.find(node_id); @@ -1987,6 +2006,10 @@ bool DistributedRDMATransport::connect_to_node(uint32_t node_id, const std::stri RDMANodeConnection conn; conn.client = std::move(client); + conn.endpoint_addr = addr; + conn.endpoint_port = port; + conn.remote_addr = remote_addr; + conn.remote_buffer_size = remote_buffer_size; conn.connected = true; connections_[node_id] = std::move(conn); @@ -2064,9 +2087,28 @@ bool DistributedRDMATransport::send_message_wait_response(uint32_t dst_node, con memcpy(rdma_req.data, req.payload.mem.data, std::min(sizeof(rdma_req.data), sizeof(req.payload.mem.data))); RDMAResponse rdma_resp; - if (it->second.client->send_request(rdma_req, rdma_resp) != 0) { - return false; + bool delivered = false; + for (int attempt = 0; attempt < 16; ++attempt) { + if (it->second.client && it->second.client->send_request(rdma_req, rdma_resp) == 0) { + delivered = true; + break; + } + + // A Soft-RoCE QP can remain nominally RTS while no completion ever + // arrives. Recreate that one directed connection and replay the + // idempotent cacheline READ/WRITE instead of wedging the whole MPI job. + SPDLOG_WARN("Reconnecting RDMA peer node {} after request failure (attempt {}/16)", dst_node, attempt + 1); + auto replacement = std::make_unique(it->second.endpoint_addr, it->second.endpoint_port); + if (replacement->connect() == 0) { + it->second.client = std::move(replacement); + it->second.connected = true; + } else { + it->second.connected = false; + std::this_thread::sleep_for(std::chrono::milliseconds(100)); + } } + if (!delivered) + return false; // Unpack response memset(&resp, 0, sizeof(resp)); @@ -2150,7 +2192,10 @@ RDMACalibrationResult DistributedRDMATransport::calibrate_node(uint32_t dst_node RDMARequest req; memset(&req, 0, sizeof(req)); req.op_type = RDMA_OP_READ; - req.addr = 0; + // Use the first valid byte of this peer's global range. Address zero is + // only valid on node 0 and caused every other peer to log 1000 failed + // backing-memory reads during calibration. + req.addr = it->second.remote_addr; req.size = 1; RDMAResponse resp; @@ -2246,9 +2291,12 @@ RDMACalibrationResult DistributedRDMATransport::get_aggregate_calibration() cons void DistributedRDMATransport::accept_loop() { while (running_) { - if (server_->accept_connection() == 0) { + auto client = server_->accept_connection(); + if (client) { SPDLOG_INFO("RDMA transport accepted incoming connection"); - server_->handle_client(); + std::lock_guard lock(incoming_mutex_); + incoming_connections_.push_back(client); + incoming_threads_.emplace_back([this, client]() { server_->handle_client(client); }); } } } @@ -2271,6 +2319,46 @@ bool DistributedMemoryServer::initialize_rdma_transport() { return false; } + // RDMAServer's fallback callback returns a synthetic success response. + // Distributed mode needs incoming READ/WRITE requests to operate on this + // node's local backing memory so that remote data is actually persistent. + rdma_transport_->set_message_handler([this](const RDMAMessage &request, RDMAMessage &response) { + memset(&response, 0, sizeof(response)); + + const RDMARequest &req = request.request; + RDMAResponse &resp = response.response; + const size_t size = std::min(req.size, RDMA_CACHELINE_SIZE); + + if (size == 0 || req.size > RDMA_CACHELINE_SIZE) { + resp.status = 1; + return; + } + + switch (req.op_type) { + case RDMA_OP_READ: + if (local_memory_->read_cacheline(req.addr, resp.data, size)) { + resp.status = 0; + resp.latency_ns = static_cast(controller_->dramlatency); + local_reads_++; + } else { + resp.status = 1; + } + break; + case RDMA_OP_WRITE: + if (local_memory_->write_cacheline(req.addr, req.data, size)) { + resp.status = 0; + resp.latency_ns = static_cast(controller_->dramlatency) + 50; + local_writes_++; + } else { + resp.status = 1; + } + break; + default: + resp.status = 1; + break; + } + }); + SPDLOG_INFO("RDMA transport initialized on {}:{}", tcp_addr_, tcp_transport_port_ + 1000); // Register with CoherencyEngine (reuse TCP transport interface via adapter) @@ -2310,13 +2398,6 @@ bool DistributedMemoryServer::connect_rdma_node(uint32_t node_id, const std::str return false; } - if (!rdma_transport_->connect_to_node(node_id, addr, port)) { - SPDLOG_ERROR("Failed to connect RDMA to node {} at {}:{}", node_id, addr, port); - return false; - } - - SPDLOG_INFO("RDMA connected to node {} at {}:{}", node_id, addr, port); - // Create virtual endpoint for the remote node uint64_t peer_capacity = memory_capacity_mb_ * 1024ULL * 1024ULL; uint64_t peer_base_addr = node_id * peer_capacity; @@ -2330,6 +2411,13 @@ bool DistributedMemoryServer::connect_rdma_node(uint32_t node_id, const std::str } } + if (!rdma_transport_->connect_to_node(node_id, addr, port, peer_base_addr, peer_capacity)) { + SPDLOG_ERROR("Failed to connect RDMA to node {} at {}:{}", node_id, addr, port); + return false; + } + + SPDLOG_INFO("RDMA connected to node {} at {}:{}", node_id, addr, port); + FabricLinkConfig link_cfg{50.0, 50.0, 64}; // 50ns hop, 50GB/s, 64 credits (RDMA advantage) auto *remote = controller_->add_remote_endpoint(node_id, peer_base_addr, peer_capacity, link_cfg); remote->msg_manager_ = msg_manager_.get(); diff --git a/src/main_server.cc b/src/main_server.cc index 4b706ae..6c6554a 100644 --- a/src/main_server.cc +++ b/src/main_server.cc @@ -899,7 +899,16 @@ int main(int argc, char *argv[]) { SPDLOG_INFO("Connecting to {} RDMA peer(s)...", tcp_peers.size()); for (const auto &peer : tcp_peers) { // RDMA port is TCP port + 1000 by convention - if (dist_server.connect_rdma_node(peer.node_id, peer.addr, peer.port + 1000)) { + bool connected = false; + for (int attempt = 1; attempt <= 20 && !connected; ++attempt) { + connected = dist_server.connect_rdma_node(peer.node_id, peer.addr, peer.port + 1000); + if (!connected && attempt < 20) { + SPDLOG_WARN("RDMA peer node {} not ready (attempt {}/20); retrying", peer.node_id, + attempt); + std::this_thread::sleep_for(std::chrono::milliseconds(500)); + } + } + if (connected) { SPDLOG_INFO("Connected to RDMA peer node {} at {}:{}", peer.node_id, peer.addr, peer.port + 1000); } else { diff --git a/src/rdma_communication.cpp b/src/rdma_communication.cpp index c0a5206..266ac4c 100644 --- a/src/rdma_communication.cpp +++ b/src/rdma_communication.cpp @@ -3,6 +3,7 @@ #include #include #include +#include #include #include @@ -97,13 +98,17 @@ int RDMAConnection::setup_qp_parameters(struct ibv_qp_init_attr &qp_attr) { int RDMAConnection::post_receive() { struct ibv_recv_wr wr, *bad_wr = nullptr; struct ibv_sge sge; + // Keep receive storage disjoint from the send WR. Reusing a buffer that is + // still owned by the receive queue eventually drives Soft-RoCE QPs into + // WR_FLUSH_ERR/RETRY_EXC_ERR under a sustained request stream. + void *recv_buffer = static_cast(conn_info_.buffer) + sizeof(RDMAMessage); memset(&wr, 0, sizeof(wr)); - wr.wr_id = reinterpret_cast(conn_info_.buffer); + wr.wr_id = reinterpret_cast(recv_buffer); wr.sg_list = &sge; wr.num_sge = 1; - sge.addr = reinterpret_cast(conn_info_.buffer); + sge.addr = reinterpret_cast(recv_buffer); sge.length = sizeof(RDMAMessage); sge.lkey = conn_info_.mr->lkey; @@ -118,17 +123,18 @@ int RDMAConnection::post_receive() { int RDMAConnection::post_send(const RDMAMessage *msg) { struct ibv_send_wr wr, *bad_wr = nullptr; struct ibv_sge sge; + void *send_buffer = conn_info_.buffer; - memcpy(conn_info_.buffer, msg, sizeof(RDMAMessage)); + memcpy(send_buffer, msg, sizeof(RDMAMessage)); memset(&wr, 0, sizeof(wr)); - wr.wr_id = reinterpret_cast(conn_info_.buffer); + wr.wr_id = reinterpret_cast(send_buffer); wr.opcode = IBV_WR_SEND; wr.sg_list = &sge; wr.num_sge = 1; wr.send_flags = IBV_SEND_SIGNALED; - sge.addr = reinterpret_cast(conn_info_.buffer); + sge.addr = reinterpret_cast(send_buffer); sge.length = sizeof(RDMAMessage); sge.lkey = conn_info_.mr->lkey; @@ -139,8 +145,18 @@ int RDMAConnection::post_send(const RDMAMessage *msg) { struct ibv_wc wc; int ne; + const auto deadline = std::chrono::steady_clock::now() + std::chrono::seconds(5); do { + if (!connected_) + return -1; ne = ibv_poll_cq(conn_info_.send_cq, 1, &wc); + if (ne == 0) + std::this_thread::yield(); + if (ne == 0 && std::chrono::steady_clock::now() >= deadline) { + std::cerr << "Send completion timed out" << std::endl; + disconnect(); + return -1; + } } while (ne == 0); if (ne < 0 || wc.status != IBV_WC_SUCCESS) { @@ -211,16 +227,47 @@ int RDMAConnection::send_message(const RDMAMessage &msg) { #endif } -int RDMAConnection::receive_message(RDMAMessage &msg) { +int RDMAConnection::receive_message(RDMAMessage &msg, int timeout_ms) { if (!connected_) return -1; #ifdef HAS_RDMA struct ibv_wc wc; int ne; + const auto deadline = timeout_ms < 0 ? std::chrono::steady_clock::time_point::max() + : std::chrono::steady_clock::now() + std::chrono::milliseconds(timeout_ms); - do { + while (connected_) { ne = ibv_poll_cq(conn_info_.recv_cq, 1, &wc); - } while (ne == 0); + if (ne != 0) + break; + + // Wait for a completion-channel event rather than burning one CPU per + // idle inbound peer. The bounded poll also lets disconnect() stop the + // handler promptly during shutdown. + struct pollfd pfd; + pfd.fd = conn_info_.comp_channel->fd; + pfd.events = POLLIN; + pfd.revents = 0; + int ready = poll(&pfd, 1, 100); + if (!connected_) + return -1; + if (std::chrono::steady_clock::now() >= deadline) { + std::cerr << "Receive completion timed out" << std::endl; + disconnect(); + return -1; + } + if (ready > 0) { + struct ibv_cq *event_cq = nullptr; + void *event_context = nullptr; + if (ibv_get_cq_event(conn_info_.comp_channel, &event_cq, &event_context) == 0) { + ibv_ack_cq_events(event_cq, 1); + if (ibv_req_notify_cq(event_cq, 0)) + return -1; + } + } else if (ready < 0) { + return -1; + } + } if (ne < 0 || wc.status != IBV_WC_SUCCESS) { std::cerr << "Receive failed with status: " << wc.status << std::endl; @@ -246,6 +293,52 @@ void RDMAConnection::disconnect() { #endif } +#ifdef HAS_RDMA +int RDMAConnection::accept_cm_id(struct rdma_cm_id *id) { + cm_id_ = id; + conn_info_.context = cm_id_->verbs; + + if (setup_connection_resources() < 0) + return -1; + + struct ibv_qp_init_attr qp_attr; + setup_qp_parameters(qp_attr); + if (rdma_create_qp(cm_id_, conn_info_.pd, &qp_attr)) { + std::cerr << "Failed to create accepted QP" << std::endl; + return -1; + } + conn_info_.qp = cm_id_->qp; + + // Requests are strictly serialized per connection. One posted receive is + // sufficient and receive_message() reposts it after each completion. + if (post_receive() < 0) + return -1; + + struct rdma_conn_param conn_param; + memset(&conn_param, 0, sizeof(conn_param)); + conn_param.initiator_depth = 1; + conn_param.responder_resources = 1; + // Soft-RoCE runs RC over UDP/IP and can see transient packet loss or an + // RNR window under synchronized load. Zero (the memset default) makes the + // QP fatal on the first retry; use the maximum retry budget instead. + conn_param.retry_count = 7; + conn_param.rnr_retry_count = 7; + if (rdma_accept(cm_id_, &conn_param)) { + std::cerr << "Failed to accept connection" << std::endl; + return -1; + } + return 0; +} +#endif + +void RDMAConnection::mark_connected() { +#ifdef HAS_RDMA + conn_info_.connected = true; +#endif + connected_ = true; + running_ = true; +} + // ---- RDMAServer ---- RDMAServer::RDMAServer(const std::string &addr, uint16_t port) : bind_addr_(addr), port_(port) { @@ -293,7 +386,9 @@ int RDMAServer::start() { return -1; } - if (rdma_listen(listen_id_, 10)) { + // A 16-node full mesh creates 15 simultaneous inbound connections per + // server. Leave headroom for connection bursts and reconnects. + if (rdma_listen(listen_id_, 64)) { std::cerr << "Failed to listen on RDMA" << std::endl; return -1; } @@ -307,79 +402,71 @@ int RDMAServer::start() { #endif } -int RDMAServer::accept_connection() { +std::shared_ptr RDMAServer::accept_connection() { #ifdef HAS_RDMA - struct rdma_cm_event *event = nullptr; - - if (rdma_get_cm_event(event_channel_, &event)) { - std::cerr << "Failed to get CM event" << std::endl; - return -1; - } - - if (event->event == RDMA_CM_EVENT_CONNECT_REQUEST) { - cm_id_ = event->id; - conn_info_.context = cm_id_->verbs; + while (running_) { + struct pollfd pfd; + pfd.fd = event_channel_->fd; + pfd.events = POLLIN; + pfd.revents = 0; + int ready = poll(&pfd, 1, 200); + if (ready == 0) + return nullptr; + if (ready < 0) { + if (running_) + std::cerr << "Failed to poll RDMA CM event channel" << std::endl; + return nullptr; + } - if (setup_connection_resources() < 0) { - rdma_ack_cm_event(event); - return -1; + struct rdma_cm_event *event = nullptr; + if (rdma_get_cm_event(event_channel_, &event)) { + if (running_) + std::cerr << "Failed to get CM event" << std::endl; + return nullptr; } - struct ibv_qp_init_attr qp_attr; - setup_qp_parameters(qp_attr); + const auto event_type = event->event; + struct rdma_cm_id *event_id = event->id; - if (rdma_create_qp(cm_id_, conn_info_.pd, &qp_attr)) { - std::cerr << "Failed to create QP" << std::endl; + if (event_type == RDMA_CM_EVENT_CONNECT_REQUEST) { + auto client = std::make_shared(); + client->set_message_handler(message_handler_); + if (client->accept_cm_id(event_id) == 0) + pending_connections_[event_id] = client; rdma_ack_cm_event(event); - return -1; + continue; } - conn_info_.qp = cm_id_->qp; - - for (int i = 0; i < RDMA_MAX_WR; i++) { - if (post_receive() < 0) { - break; + if (event_type == RDMA_CM_EVENT_ESTABLISHED) { + auto it = pending_connections_.find(event_id); + if (it != pending_connections_.end()) { + auto client = it->second; + pending_connections_.erase(it); + client->mark_connected(); + rdma_ack_cm_event(event); + std::cout << "RDMA connection established" << std::endl; + return client; } } - struct rdma_conn_param conn_param; - memset(&conn_param, 0, sizeof(conn_param)); - conn_param.initiator_depth = 1; - conn_param.responder_resources = 1; - - if (rdma_accept(cm_id_, &conn_param)) { - std::cerr << "Failed to accept connection" << std::endl; - rdma_ack_cm_event(event); - return -1; - } - } - - rdma_ack_cm_event(event); - - if (rdma_get_cm_event(event_channel_, &event)) { - return -1; - } + if (event_type == RDMA_CM_EVENT_REJECTED || event_type == RDMA_CM_EVENT_DISCONNECTED || + event_type == RDMA_CM_EVENT_CONNECT_ERROR) + pending_connections_.erase(event_id); - if (event->event == RDMA_CM_EVENT_ESTABLISHED) { - conn_info_.connected = true; - connected_ = true; - std::cout << "RDMA connection established" << std::endl; + rdma_ack_cm_event(event); } - - rdma_ack_cm_event(event); - return 0; + return nullptr; #else - return -1; + return nullptr; #endif } -void RDMAServer::handle_client() { - while (running_ && connected_) { +void RDMAServer::handle_client(const std::shared_ptr &client) { + while (running_ && client && client->is_connected()) { RDMAMessage recv_msg, send_msg; - if (receive_message(recv_msg) < 0) { + if (client->receive_message(recv_msg) < 0) break; - } if (message_handler_) { message_handler_(recv_msg, send_msg); @@ -388,12 +475,12 @@ void RDMAServer::handle_client() { send_msg.response.latency_ns = 100; } - if (send_message(send_msg) < 0) { + if (client->send_message(send_msg) < 0) break; - } } - connected_ = false; + if (client) + client->disconnect(); } void RDMAServer::stop() { @@ -461,11 +548,11 @@ int RDMAClient::connect() { conn_info_.qp = cm_id_->qp; - for (int i = 0; i < RDMA_MAX_WR; i++) { - if (post_receive() < 0) { - break; - } - } + // Requests are serialized per peer, and every completion reposts the + // receive. Posting the same buffer hundreds of times is invalid and can + // race a response send against another receive completion. + if (post_receive() < 0) + return -1; if (rdma_resolve_route(cm_id_, 2000)) { std::cerr << "Failed to resolve route" << std::endl; @@ -488,6 +575,8 @@ int RDMAClient::connect() { memset(&conn_param, 0, sizeof(conn_param)); conn_param.initiator_depth = 1; conn_param.responder_resources = 1; + conn_param.retry_count = 7; + conn_param.rnr_retry_count = 7; if (rdma_connect(cm_id_, &conn_param)) { std::cerr << "Failed to connect" << std::endl; @@ -530,7 +619,7 @@ int RDMAClient::send_request(const RDMARequest &req, RDMAResponse &resp) { return -1; } - if (receive_message(msg) < 0) { + if (receive_message(msg, 5000) < 0) { return -1; } From d94d9b69ff2d5ef7470e891f9a1d470617be658b Mon Sep 17 00:00:00 2001 From: victoryang00 Date: Fri, 7 Aug 2026 23:09:20 -0700 Subject: [PATCH 2/4] Batch distributed RDMA range transfers Signed-off-by: victoryang00 --- include/distributed_server.h | 4 +- include/rdma_communication.h | 6 +- include/shared_memory_manager.h | 2 + src/distributed_server.cpp | 145 ++++++++++++++++++++++++++++++-- src/rdma_communication.cpp | 5 +- src/shared_memory_manager.cc | 59 +++++++++++-- 6 files changed, 204 insertions(+), 17 deletions(-) diff --git a/include/distributed_server.h b/include/distributed_server.h index 8b159cf..c04a0df 100644 --- a/include/distributed_server.h +++ b/include/distributed_server.h @@ -316,8 +316,8 @@ typedef struct { #include /* TCP transport support */ -#include "tcp_communication.h" #include "shared_memory_manager.h" +#include "tcp_communication.h" /* RDMA transport support */ #include "rdma_communication.h" @@ -776,6 +776,8 @@ class DistributedMemoryServer { /* Memory operations (may be forwarded to remote nodes) */ int read(uint64_t addr, void *data, size_t size, uint64_t *latency_ns); int write(uint64_t addr, const void *data, size_t size, uint64_t *latency_ns); + int read_bulk(uint64_t addr, void *data, size_t size, uint64_t *latency_ns); + int write_bulk(uint64_t addr, const void *data, size_t size, uint64_t *latency_ns); int atomic_faa(uint64_t addr, uint64_t value, uint64_t *old_value); int atomic_cas(uint64_t addr, uint64_t expected, uint64_t desired, uint64_t *old_value); void fence(); diff --git a/include/rdma_communication.h b/include/rdma_communication.h index ef22441..7fb5682 100644 --- a/include/rdma_communication.h +++ b/include/rdma_communication.h @@ -13,7 +13,7 @@ #include #endif -#define RDMA_BUFFER_SIZE 4096 +#define RDMA_BUFFER_SIZE (64 * 1024) #define RDMA_CQ_SIZE 1024 #define RDMA_MAX_WR 512 #define RDMA_CACHELINE_SIZE 64 @@ -27,14 +27,14 @@ struct RDMARequest { uint64_t timestamp; uint8_t host_id; uint64_t virtual_addr; - uint8_t data[RDMA_CACHELINE_SIZE]; + uint8_t data[RDMA_BUFFER_SIZE]; } __attribute__((packed)); struct RDMAResponse { uint8_t status; uint64_t latency_ns; uint8_t cache_state; - uint8_t data[RDMA_CACHELINE_SIZE]; + uint8_t data[RDMA_BUFFER_SIZE]; } __attribute__((packed)); struct RDMAMessage { diff --git a/include/shared_memory_manager.h b/include/shared_memory_manager.h index e48787b..4da59b6 100644 --- a/include/shared_memory_manager.h +++ b/include/shared_memory_manager.h @@ -143,6 +143,8 @@ class SharedMemoryManager { uint8_t *get_cacheline_data(uint64_t cacheline_addr); bool read_cacheline(uint64_t addr, uint8_t *buffer, size_t size); bool write_cacheline(uint64_t addr, const uint8_t *data, size_t size); + bool read_range(uint64_t addr, uint8_t *buffer, size_t size); + bool write_range(uint64_t addr, const uint8_t *data, size_t size); bool atomic_fetch_add_uint64(uint64_t addr, uint64_t value, uint64_t *old_value); bool atomic_compare_exchange_uint64(uint64_t addr, uint64_t expected, uint64_t desired, uint64_t *old_value); bool flush(); diff --git a/src/distributed_server.cpp b/src/distributed_server.cpp index 286241a..374f492 100644 --- a/src/distributed_server.cpp +++ b/src/distributed_server.cpp @@ -57,9 +57,40 @@ static constexpr uint8_t DIST_OP_ATOMIC_CAS = 4; static constexpr uint8_t DIST_OP_FENCE = 5; static constexpr uint8_t DIST_OP_LSA_READ = 6; static constexpr uint8_t DIST_OP_LSA_WRITE = 7; +static constexpr uint8_t DIST_OP_BULK_READ = 8; +static constexpr uint8_t DIST_OP_BULK_WRITE = 9; +static constexpr uint64_t DIST_MAX_BULK_SIZE = 64ULL * 1024 * 1024; namespace { +bool recv_full_fd(int fd, void *buffer, size_t size) { + auto *cursor = static_cast(buffer); + while (size > 0) { + ssize_t received = recv(fd, cursor, size, 0); + if (received < 0 && errno == EINTR) + continue; + if (received <= 0) + return false; + cursor += received; + size -= static_cast(received); + } + return true; +} + +bool send_full_fd(int fd, const void *buffer, size_t size) { + const auto *cursor = static_cast(buffer); + while (size > 0) { + ssize_t sent = send(fd, cursor, size, MSG_NOSIGNAL); + if (sent < 0 && errno == EINTR) + continue; + if (sent <= 0) + return false; + cursor += sent; + size -= static_cast(sent); + } + return true; +} + std::string expand_node_backing_path(std::string path, uint32_t node_id) { constexpr const char *kNodeToken = "{node}"; size_t pos = path.find(kNodeToken); @@ -905,7 +936,7 @@ void DistributedMemoryServer::handle_atomic_request(const dist_message_t &req, d ok = local_memory_->atomic_fetch_add_uint64(addr, req.payload.mem.value, &old_value); } else { ok = local_memory_->atomic_compare_exchange_uint64(addr, req.payload.mem.expected, req.payload.mem.value, - &old_value); + &old_value); } if (!ok) { @@ -1081,6 +1112,72 @@ int DistributedMemoryServer::write(uint64_t addr, const void *data, size_t size, return forward_write(target_node, addr, data, size, latency_ns); } +int DistributedMemoryServer::read_bulk(uint64_t addr, void *data, size_t size, uint64_t *latency_ns) { + if (!data || size == 0 || addr > UINT64_MAX - (size - 1)) + return -1; + const uint32_t target_node = get_node_for_address(addr); + if (get_node_for_address(addr + size - 1) != target_node) + return -1; + + if (target_node == node_id_ || target_node == UINT32_MAX) { + if (!local_memory_->read_range(addr, static_cast(data), size)) + return -1; + local_reads_++; + *latency_ns = static_cast(controller_->dramlatency); + return 0; + } + if (transport_mode_ != DistTransportMode::RDMA || !rdma_transport_ || !rdma_transport_->is_connected(target_node)) + return -1; + + auto *cursor = static_cast(data); + uint64_t total_latency = 0; + for (size_t offset = 0; offset < size; offset += RDMA_BUFFER_SIZE) { + const size_t chunk = std::min(RDMA_BUFFER_SIZE, size - offset); + const auto start = std::chrono::steady_clock::now(); + if (!rdma_transport_->rdma_read(target_node, addr + offset, cursor + offset, chunk)) + return -1; + total_latency += + std::chrono::duration_cast(std::chrono::steady_clock::now() - start).count(); + remote_reads_++; + forwarded_requests_++; + } + *latency_ns = total_latency; + return 0; +} + +int DistributedMemoryServer::write_bulk(uint64_t addr, const void *data, size_t size, uint64_t *latency_ns) { + if (!data || size == 0 || addr > UINT64_MAX - (size - 1)) + return -1; + const uint32_t target_node = get_node_for_address(addr); + if (get_node_for_address(addr + size - 1) != target_node) + return -1; + + if (target_node == node_id_ || target_node == UINT32_MAX) { + if (!local_memory_->write_range(addr, static_cast(data), size)) + return -1; + local_writes_++; + *latency_ns = static_cast(controller_->dramlatency) + 50; + return 0; + } + if (transport_mode_ != DistTransportMode::RDMA || !rdma_transport_ || !rdma_transport_->is_connected(target_node)) + return -1; + + const auto *cursor = static_cast(data); + uint64_t total_latency = 0; + for (size_t offset = 0; offset < size; offset += RDMA_BUFFER_SIZE) { + const size_t chunk = std::min(RDMA_BUFFER_SIZE, size - offset); + const auto start = std::chrono::steady_clock::now(); + if (!rdma_transport_->rdma_write(target_node, addr + offset, cursor + offset, chunk)) + return -1; + total_latency += + std::chrono::duration_cast(std::chrono::steady_clock::now() - start).count(); + remote_writes_++; + forwarded_requests_++; + } + *latency_ns = total_latency; + return 0; +} + int DistributedMemoryServer::forward_read(uint32_t target_node, uint64_t addr, void *data, size_t size, uint64_t *latency_ns) { forwarded_requests_++; @@ -1990,7 +2087,7 @@ void DistributedRDMATransport::shutdown() { } bool DistributedRDMATransport::connect_to_node(uint32_t node_id, const std::string &addr, uint16_t port, - uint64_t remote_addr, size_t remote_buffer_size) { + uint64_t remote_addr, size_t remote_buffer_size) { std::lock_guard lock(connections_mutex_); auto it = connections_.find(node_id); @@ -2327,16 +2424,16 @@ bool DistributedMemoryServer::initialize_rdma_transport() { const RDMARequest &req = request.request; RDMAResponse &resp = response.response; - const size_t size = std::min(req.size, RDMA_CACHELINE_SIZE); + const size_t size = std::min(req.size, RDMA_BUFFER_SIZE); - if (size == 0 || req.size > RDMA_CACHELINE_SIZE) { + if (size == 0 || req.size > RDMA_BUFFER_SIZE) { resp.status = 1; return; } switch (req.op_type) { case RDMA_OP_READ: - if (local_memory_->read_cacheline(req.addr, resp.data, size)) { + if (local_memory_->read_range(req.addr, resp.data, size)) { resp.status = 0; resp.latency_ns = static_cast(controller_->dramlatency); local_reads_++; @@ -2345,7 +2442,7 @@ bool DistributedMemoryServer::initialize_rdma_transport() { } break; case RDMA_OP_WRITE: - if (local_memory_->write_cacheline(req.addr, req.data, size)) { + if (local_memory_->write_range(req.addr, req.data, size)) { resp.status = 0; resp.latency_ns = static_cast(controller_->dramlatency) + 50; local_writes_++; @@ -2926,6 +3023,7 @@ void DistributedMemoryServer::handle_tcp_client(int client_fd, int client_id) { DistServerResponse resp; memset(&resp, 0, sizeof(resp)); + std::vector bulk_response; // Handle the request using distributed read/write (with proper forwarding) switch (req.op_type) { @@ -2989,6 +3087,38 @@ void DistributedMemoryServer::handle_tcp_client(int client_fd, int client_id) { break; } + case DIST_OP_BULK_READ: { + if (req.size == 0 || req.size > DIST_MAX_BULK_SIZE || req.addr > UINT64_MAX - (req.size - 1)) { + resp.status = 1; + break; + } + bulk_response.resize(static_cast(req.size)); + uint64_t latency_ns = 0; + int ret = read_bulk(req.addr, bulk_response.data(), bulk_response.size(), &latency_ns); + resp.status = ret == 0 ? 0 : 1; + resp.latency_ns = latency_ns; + if (ret != 0) + bulk_response.clear(); + break; + } + + case DIST_OP_BULK_WRITE: { + if (req.size == 0 || req.size > DIST_MAX_BULK_SIZE || req.addr > UINT64_MAX - (req.size - 1)) { + resp.status = 1; + break; + } + std::vector payload(static_cast(req.size)); + if (!recv_full_fd(client_fd, payload.data(), payload.size())) { + resp.status = 1; + break; + } + uint64_t latency_ns = 0; + int ret = write_bulk(req.addr, payload.data(), payload.size(), &latency_ns); + resp.status = ret == 0 ? 0 : 1; + resp.latency_ns = latency_ns; + break; + } + case DIST_OP_GET_SHM_INFO: { // Return shared memory info auto shm_info = local_memory_->get_shm_info(); @@ -3007,7 +3137,8 @@ void DistributedMemoryServer::handle_tcp_client(int client_fd, int client_id) { } // Send response - if (send(client_fd, &resp, sizeof(resp), MSG_NOSIGNAL) != sizeof(resp)) { + if (!send_full_fd(client_fd, &resp, sizeof(resp)) || + (!bulk_response.empty() && !send_full_fd(client_fd, bulk_response.data(), bulk_response.size()))) { if (running_) { SPDLOG_ERROR("Node {} client {}: Failed to send response: {}", node_id_, client_id, strerror(errno)); } diff --git a/src/rdma_communication.cpp b/src/rdma_communication.cpp index 266ac4c..05ac423 100644 --- a/src/rdma_communication.cpp +++ b/src/rdma_communication.cpp @@ -59,7 +59,10 @@ int RDMAConnection::setup_connection_resources() { } int RDMAConnection::register_memory_region() { - conn_info_.buffer_size = RDMA_BUFFER_SIZE * sizeof(RDMAMessage); + // Requests are serialized and use one send plus one receive buffer. The + // old allocation multiplied the message size by RDMA_BUFFER_SIZE even + // though only two slots were ever addressed. + conn_info_.buffer_size = 2 * sizeof(RDMAMessage); conn_info_.buffer = malloc(conn_info_.buffer_size); if (!conn_info_.buffer) { std::cerr << "Failed to allocate buffer" << std::endl; diff --git a/src/shared_memory_manager.cc b/src/shared_memory_manager.cc index a1db10b..1f6faee 100644 --- a/src/shared_memory_manager.cc +++ b/src/shared_memory_manager.cc @@ -11,10 +11,10 @@ #include #include #include -#include #include #include #include +#include #include #ifdef CXLMEMSIM_HAS_SSD_STREAMING_BACKEND @@ -224,9 +224,9 @@ bool SharedMemoryManager::create_ssd_streaming_backend() { SPDLOG_INFO("SSD streaming backend initialized: path={} capacity={} page={} chunk={} cache_pages={} " "read_ahead_pages={} io_uring={} odirect={}", - ssd_config.backing_path, ssd_config.capacity_bytes, ssd_config.page_size, - ssd_config.io_chunk_size, ssd_config.cache_pages, ssd_config.read_ahead_pages, - ssd_config.use_io_uring, ssd_config.use_odirect); + ssd_config.backing_path, ssd_config.capacity_bytes, ssd_config.page_size, ssd_config.io_chunk_size, + ssd_config.cache_pages, ssd_config.read_ahead_pages, ssd_config.use_io_uring, + ssd_config.use_odirect); return true; } catch (const std::exception &e) { SPDLOG_ERROR("Exception while initializing SSD streaming backend: {}", e.what()); @@ -607,6 +607,55 @@ bool SharedMemoryManager::write_cacheline(uint64_t addr, const uint8_t *data, si return true; } +bool SharedMemoryManager::read_range(uint64_t addr, uint8_t *buffer, size_t size) { + if (!header || !buffer || size == 0) + return false; + const uint64_t capacity = header->num_cachelines * SHM_CACHELINE_SIZE; + const uint64_t base = header->base_addr; + if (capacity == 0 || (base != 0 && addr < base)) + return false; + const uint64_t offset = base == 0 ? addr % capacity : addr - base; + if (offset > capacity || size > capacity - offset) + return false; + + if (backing_mode == BackingMode::SsdStream) { +#ifdef CXLMEMSIM_HAS_SSD_STREAMING_BACKEND + return ssd_backend && call_backend_bool([&]() { return ssd_backend->read(offset, buffer, size); }); +#else + return false; +#endif + } + if (!data_area) + return false; + memcpy(buffer, data_area + offset, size); + return true; +} + +bool SharedMemoryManager::write_range(uint64_t addr, const uint8_t *data, size_t size) { + if (!header || !data || size == 0) + return false; + const uint64_t capacity = header->num_cachelines * SHM_CACHELINE_SIZE; + const uint64_t base = header->base_addr; + if (capacity == 0 || (base != 0 && addr < base)) + return false; + const uint64_t offset = base == 0 ? addr % capacity : addr - base; + if (offset > capacity || size > capacity - offset) + return false; + + if (backing_mode == BackingMode::SsdStream) { +#ifdef CXLMEMSIM_HAS_SSD_STREAMING_BACKEND + return ssd_backend && call_backend_bool([&]() { return ssd_backend->write(offset, data, size); }); +#else + return false; +#endif + } + if (!data_area) + return false; + memcpy(data_area + offset, data, size); + __atomic_thread_fence(__ATOMIC_RELEASE); + return true; +} + bool SharedMemoryManager::atomic_fetch_add_uint64(uint64_t addr, uint64_t value, uint64_t *old_value) { if (!old_value) { return false; @@ -640,7 +689,7 @@ bool SharedMemoryManager::atomic_fetch_add_uint64(uint64_t addr, uint64_t value, } bool SharedMemoryManager::atomic_compare_exchange_uint64(uint64_t addr, uint64_t expected, uint64_t desired, - uint64_t *old_value) { + uint64_t *old_value) { if (!old_value) { return false; } From d4f4d11b2433ed4b54e97312844c838fcf1a71a3 Mon Sep 17 00:00:00 2001 From: victoryang00 Date: Sat, 8 Aug 2026 12:26:06 -0700 Subject: [PATCH 3/4] Bound LSA growth from client requests Signed-off-by: victoryang00 --- src/distributed_server.cpp | 18 ++++++++++++------ 1 file changed, 12 insertions(+), 6 deletions(-) diff --git a/src/distributed_server.cpp b/src/distributed_server.cpp index 374f492..a9f8ac4 100644 --- a/src/distributed_server.cpp +++ b/src/distributed_server.cpp @@ -60,6 +60,7 @@ static constexpr uint8_t DIST_OP_LSA_WRITE = 7; static constexpr uint8_t DIST_OP_BULK_READ = 8; static constexpr uint8_t DIST_OP_BULK_WRITE = 9; static constexpr uint64_t DIST_MAX_BULK_SIZE = 64ULL * 1024 * 1024; +static constexpr uint64_t DIST_MAX_LSA_SIZE = 64ULL * 1024 * 1024; namespace { @@ -1324,7 +1325,7 @@ void DistributedMemoryServer::fence() { int DistributedMemoryServer::lsa_read(uint64_t offset, void *data, size_t size) { std::lock_guard lock(lsa_mutex_); - if (offset + size > lsa_size_) { + if (!data || size == 0 || offset > lsa_size_ || size > lsa_size_ - offset) { SPDLOG_ERROR("Node {} LSA read out of bounds: offset=0x{:x} size={} lsa_size={}", node_id_, offset, size, lsa_size_); return -1; @@ -1335,12 +1336,17 @@ int DistributedMemoryServer::lsa_read(uint64_t offset, void *data, size_t size) int DistributedMemoryServer::lsa_write(uint64_t offset, const void *data, size_t size) { std::lock_guard lock(lsa_mutex_); - if (offset + size > lsa_size_) { + if (!data || size == 0 || offset > DIST_MAX_LSA_SIZE || size > DIST_MAX_LSA_SIZE - offset) { + SPDLOG_WARN("Node {} rejected LSA write: offset=0x{:x} size={} max={}", node_id_, offset, size, + DIST_MAX_LSA_SIZE); + return -1; + } + const size_t required_size = static_cast(offset + size); + if (required_size > lsa_size_) { /* Auto-grow LSA if needed */ - size_t new_size = offset + size; - SPDLOG_INFO("Node {} growing LSA from {} to {} bytes", node_id_, lsa_size_, new_size); - lsa_data_.resize(new_size, 0); - lsa_size_ = new_size; + SPDLOG_INFO("Node {} growing LSA from {} to {} bytes", node_id_, lsa_size_, required_size); + lsa_data_.resize(required_size, 0); + lsa_size_ = required_size; } memcpy(lsa_data_.data() + offset, data, size); return 0; From 05ca2542cdb7b35e5f97e1e09bfc5de703c3bfad Mon Sep 17 00:00:00 2001 From: victoryang00 Date: Sat, 8 Aug 2026 12:37:11 -0700 Subject: [PATCH 4/4] Allow production-sized bounded LSA pools Signed-off-by: victoryang00 --- src/distributed_server.cpp | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/src/distributed_server.cpp b/src/distributed_server.cpp index a9f8ac4..edc8fc0 100644 --- a/src/distributed_server.cpp +++ b/src/distributed_server.cpp @@ -60,7 +60,10 @@ static constexpr uint8_t DIST_OP_LSA_WRITE = 7; static constexpr uint8_t DIST_OP_BULK_READ = 8; static constexpr uint8_t DIST_OP_BULK_WRITE = 9; static constexpr uint64_t DIST_MAX_BULK_SIZE = 64ULL * 1024 * 1024; -static constexpr uint64_t DIST_MAX_LSA_SIZE = 64ULL * 1024 * 1024; +// A vLLM worker may address the full configured per-rank offload pool (255 MiB +// in the 70B benchmark). Keep a finite guard against hostile sparse writes, +// while leaving headroom for legitimate distributed-memory workloads. +static constexpr uint64_t DIST_MAX_LSA_SIZE = 1ULL * 1024 * 1024 * 1024; namespace {