[TE][1/N] add RPC based data transfer (#1104)

* feat: Add coro RPC communicator core implementation

- Add cororpc_communicator and cororpc_interface headers and implementation
- Implement async RPC communication layer using yalantinglibs coro_rpc
- Add CMake configuration for coro_rpc_connector module
- Add bandwidth test script for performance validation
- Support both TCP and RDMA transports via MC_RPC_PROTOCOL env var

This introduces a new transport layer based on coro_rpc for
high-performance asynchronous communication in the transfer engine.
The implementation uses zero-copy techniques via attachments and
provides both synchronous and asynchronous APIs.

* fixed issues: 1.moved config to config.h 2. renamed config to RPCCommunicator config 3. removed comment lines

* fixed the names

* fixed compilation errors

* fixed some compilation errors

* 1. updated names 2. updated CI

* reformat the code

* 1.removed useless pybind dependencies 2. renamed test file

* fixed AI assisted bugs

* fixed compile errors

* fixed cmakelists.txt

* fixed some bugs

* removed useless files

* fixed hardcoding of cmakelists.txt

* updated cmakelists.txt

* updated pybind

* update pybind

* update cmakelists.txt pybind11

* set cmakelists.txt

---------

Co-authored-by: lukotong-7 <shicanwei.scw@alibaba-inc.com>

.github/workflows/ci.yml

@@ -239,6 +239,16 @@ jobs:
         python scripts/test_tensor_api.py -n 1
       shell: bash
 
+    - name: Run RPC Communicator Bandwidth Test
+      run: |
+        source test_env/bin/activate
+        python mooncake-transfer-engine/tests/rpc_communicator_test.py server --url 127.0.0.1:9004 --data-size 1 &
+        SERVER_PID=$!
+        sleep 5 
+        timeout 10 python mooncake-transfer-engine/tests/rpc_communicator_test.py client --url 127.0.0.1:9004 --threads 2 --data-size 1 || true
+        kill $SERVER_PID 2>/dev/null || true
+        wait $SERVER_PID 2>/dev/null || true
+        
     - name: Test Mooncake EP Backend (CPU Only)
       run: |
         source test_env/bin/activate

mooncake-integration/CMakeLists.txt

@@ -36,6 +36,14 @@ if (WITH_TE)
     set_target_properties(engine PROPERTIES
         INSTALL_RPATH "$ORIGIN" 
     )
+    
+    target_link_libraries(engine PRIVATE 
+        $<TARGET_OBJECTS:rpc_communicator>
+    )
+
+    find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
+    target_link_libraries(engine PRIVATE ${Python3_LIBRARIES})
+    target_include_directories(engine PRIVATE ${Python3_INCLUDE_DIRS})
     if (USE_ASCEND_DIRECT)
         target_link_libraries(engine PUBLIC
             ascendcl

mooncake-integration/transfer_engine/transfer_engine_py.cpp

@@ -19,6 +19,7 @@
 #include <fstream>
 
 #include <pybind11/stl.h>
+#include "transport/rpc_communicator/rpc_interface.h"
 
 #ifdef USE_MNNVL
 #include "transport/nvlink_transport/nvlink_transport.h"
@@ -775,4 +776,7 @@ PYBIND11_MODULE(engine, m) {
 
     py::class_<TransferEngine, std::shared_ptr<TransferEngine>>(
         m, "InnerTransferEngine");
+
+    // Bind RpcInterface
+    mooncake::bind_rpc_interface(m);
 }

mooncake-transfer-engine/include/config.h

@@ -60,6 +60,13 @@ struct GlobalConfig {
     EndpointStoreType endpoint_store_type = EndpointStoreType::SIEVE;
 };
 
+struct RpcCommunicatorConfig {
+    std::string listen_address;
+    size_t thread_count = 0;
+    size_t timeout_seconds = 30;
+    size_t pool_size = 10;
+};
+
 void loadGlobalConfig(GlobalConfig &config);
 
 void dumpGlobalConfig();

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

@@ -0,0 +1,81 @@
+#pragma once
+
+#include <memory>
+#include <string>
+#include <vector>
+#include <future>
+#include <unordered_map>
+#include <functional>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#include <ylt/coro_rpc/coro_rpc_client.hpp>
+#include <ylt/coro_rpc/coro_rpc_server.hpp>
+#include <ylt/coro_io/client_pool.hpp>
+#include <ylt/coro_io/coro_io.hpp>
+#include <async_simple/coro/Lazy.h>
+#include "config.h"
+
+namespace mooncake {
+
+struct TensorInfo {
+    void* data_ptr = nullptr;
+    std::vector<size_t> shape;
+    std::string dtype;
+    size_t total_bytes = 0;
+};
+
+struct RpcResult {
+    int code = 0;
+    std::string err_msg;
+};
+
+class RpcCommunicator {
+   public:
+    RpcCommunicator();
+    ~RpcCommunicator();
+
+    bool initialize(const RpcCommunicatorConfig& config);
+    bool startServerImpl(bool is_async = true);
+    bool startServer();
+    bool startServerAsync();
+    void stopServer();
+
+    int sendData(const std::string& target_address, const void* data,
+                 size_t data_size);
+    async_simple::coro::Lazy<RpcResult> sendDataAsync(
+        const std::string& target_address, const void* data, size_t data_size);
+
+    int sendTensor(const std::string& target_address,
+                   const pybind11::object& tensor);
+    async_simple::coro::Lazy<int> sendTensorAsync(
+        const std::string& target_address, const TensorInfo& tensor);
+
+    int receiveData(const std::string& source_address, void* buffer,
+                    size_t buffer_size, int timeout_ms = -1);
+    async_simple::coro::Lazy<std::string> receiveDataAsync(
+        const std::string& source_address, int timeout_ms = -1);
+
+    void setDataReceiveCallback(
+        std::function<void(std::string_view, std::string_view)> callback);
+
+   private:
+    // Handler methods for RPC calls
+    void handleDataTransfer(coro_rpc::context<void> context,
+                            std::string_view data);
+    void handleTensorTransfer(coro_rpc::context<void> context);
+    void handleDataTransferWithAttachment(coro_rpc::context<void> context,
+                                          std::string_view data);
+    void handleTensorTransferWithAttachment(coro_rpc::context<void> context);
+
+    // Implementation members
+    RpcCommunicatorConfig config_;
+    bool is_server_started_ = false;
+    std::unique_ptr<coro_rpc::coro_rpc_server> server_;
+    std::function<void(std::string_view, std::string_view)>
+        data_receive_callback_;
+    pybind11::handle py_callback_;
+    std::shared_ptr<coro_io::client_pools<coro_rpc::coro_rpc_client>>
+        client_pools_;
+};
+
+}  // namespace mooncake

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

@@ -0,0 +1,91 @@
+#pragma once
+
+#include <memory>
+#include <string>
+#include <vector>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+
+namespace mooncake {
+
+class RpcInterface {
+   public:
+    struct ReceivedData {
+        std::string source_address;
+        std::string data;
+        size_t data_size = 0;
+
+        pybind11::bytes getBytes() const { return pybind11::bytes(data); }
+        pybind11::memoryview getMemoryView() const {
+            return pybind11::memoryview::from_memory(
+                const_cast<void*>(static_cast<const void*>(data.data())),
+                data.size(), true);
+        }
+    };
+
+    struct ReceivedTensor {
+        std::string source_address;
+        std::string data;
+        std::vector<size_t> shape;
+        std::string dtype;
+        size_t total_bytes = 0;
+        size_t getDataSize() const { return data.size(); }
+        pybind11::bytes getDataAsBytes() const { return pybind11::bytes(data); }
+        pybind11::memoryview getMemoryView() const {
+            return pybind11::memoryview::from_memory(
+                const_cast<void*>(static_cast<const void*>(data.data())),
+                data.size(), true);
+        }
+    };
+
+    class Impl;
+
+    RpcInterface();
+    ~RpcInterface();
+
+    bool initialize(const std::string& listen_address = "",
+                    size_t thread_count = 0, size_t timeout_seconds = 30,
+                    size_t pool_size = 10);
+
+    // Convenience methods for common use cases
+    bool initializeClient(size_t pool_size = 10, size_t timeout_seconds = 30);
+    bool initializeServer(const std::string& listen_address,
+                          size_t thread_count = 8, size_t timeout_seconds = 30);
+
+    bool startServer();
+    bool startServerAsync();
+    bool startServerImpl(bool is_async = true);
+    void stopServer();
+
+    int sendData(const std::string& target_address, pybind11::handle data);
+    pybind11::object sendDataAsync(const std::string& target_address,
+                                   pybind11::handle data,
+                                   pybind11::handle loop);
+
+    int sendTensor(const std::string& target_address, pybind11::handle tensor);
+    pybind11::object sendTensorAsync(const std::string& target_address,
+                                     pybind11::handle tensor,
+                                     pybind11::handle loop);
+
+    void setDataReceiveCallback(pybind11::function callback);
+    void setTensorReceiveCallback(pybind11::function callback);
+
+    void handleIncomingData(std::string_view source_address,
+                            std::string_view data);
+    void handleIncomingTensor(std::string_view source_address,
+                              std::string_view data,
+                              const std::vector<size_t>& shape,
+                              std::string_view dtype);
+
+   private:
+    std::unique_ptr<Impl> impl_;
+};
+
+std::unique_ptr<RpcInterface> createRpcClient(uint64_t local_rank = 0,
+                                              uint64_t world_size = 1);
+std::unique_ptr<RpcInterface> createRpcServer(uint64_t local_rank = 0,
+                                              uint64_t world_size = 1);
+
+void bind_rpc_interface(pybind11::module_& m);
+
+}  // namespace mooncake

mooncake-transfer-engine/src/transport/CMakeLists.txt

@@ -1,7 +1,8 @@
 file(GLOB XPORT_SOURCES "*.cpp")
 
 add_subdirectory(rdma_transport)
-add_library(transport OBJECT ${XPORT_SOURCES} $<TARGET_OBJECTS:rdma_transport>)
+add_subdirectory(rpc_communicator)
+add_library(transport OBJECT ${XPORT_SOURCES} $<TARGET_OBJECTS:rdma_transport> $<TARGET_OBJECTS:rpc_communicator>)
 target_link_libraries(transport PRIVATE JsonCpp::JsonCpp yalantinglibs::yalantinglibs glog::glog pthread)
 
 if (USE_TCP)

mooncake-transfer-engine/src/transport/rpc_communicator/CMakeLists.txt

@@ -0,0 +1,41 @@
+file(GLOB CORO_RPC_SOURCES "*.cpp")
+
+find_package(Python3 COMPONENTS Interpreter Development REQUIRED)
+
+add_library(rpc_communicator OBJECT ${CORO_RPC_SOURCES})
+
+target_link_libraries(rpc_communicator 
+    PRIVATE 
+    JsonCpp::JsonCpp 
+    yalantinglibs::yalantinglibs 
+    glog::glog 
+    pthread 
+    ${Python3_LIBRARIES}
+)
+
+target_include_directories(rpc_communicator 
+    PRIVATE 
+    ${Python3_INCLUDE_DIRS}
+)
+
+if(TARGET pybind11::headers)
+    target_link_libraries(rpc_communicator PRIVATE pybind11::headers)
+elseif(DEFINED pybind11_INCLUDE_DIR)
+    target_include_directories(rpc_communicator PRIVATE ${pybind11_INCLUDE_DIR})
+else()
+    set(_PYBIND11_DIR "")
+    if(EXISTS "${CMAKE_SOURCE_DIR}/extern/pybind11/include/pybind11/pybind11.h")
+        set(_PYBIND11_DIR "${CMAKE_SOURCE_DIR}/extern/pybind11/include")
+    else()
+        get_filename_component(_PROJECT_ROOT "${CMAKE_CURRENT_SOURCE_DIR}/../../../../" ABSOLUTE)
+        set(_PYBIND11_DIR "${_PROJECT_ROOT}/extern/pybind11/include")
+        if(NOT EXISTS "${_PYBIND11_DIR}/pybind11/pybind11.h")
+            set(_PYBIND11_DIR "")
+        endif()
+    endif()
+    if(_PYBIND11_DIR)
+        target_include_directories(rpc_communicator PRIVATE "${_PYBIND11_DIR}")
+    else()
+        message(FATAL_ERROR "Could not find pybind11")
+    endif()
+endif()

mooncake-transfer-engine/src/transport/rpc_communicator/rpc_communicator.cpp

@@ -0,0 +1,419 @@
+#include "transport/rpc_communicator/rpc_communicator.h"
+#include <iostream>
+#include <thread>
+#include <functional>
+#include <glog/logging.h>
+#include <ylt/coro_rpc/coro_rpc_client.hpp>
+#include <ylt/coro_rpc/coro_rpc_server.hpp>
+#include <ylt/coro_io/client_pool.hpp>
+#include <ylt/coro_io/coro_io.hpp>
+#include <pybind11/pybind11.h>
+#include <pybind11/pytypes.h>
+#include "async_simple/coro/SyncAwait.h"
+
+namespace py = pybind11;
+
+namespace mooncake {
+
+class py_rpc_context {
+   public:
+    void response_msg(py::buffer msg, py::object done) {
+        py::buffer_info info = msg.request();
+        const char* data = static_cast<char*>(info.ptr);
+        context_.get_context_info()->set_response_attachment(
+            std::string_view(data, info.size));
+        context_.get_context_info()->set_complete_handler(
+            [done](const std::error_code& ec, std::size_t) {
+                py::gil_scoped_acquire acquire;
+                done(!ec);
+            });
+        context_.response_msg();
+    }
+
+    coro_rpc::context<void> context_;
+};
+
+RpcCommunicator::RpcCommunicator() {}
+
+RpcCommunicator::~RpcCommunicator() { stopServer(); }
+
+void RpcCommunicator::setDataReceiveCallback(
+    std::function<void(std::string_view, std::string_view)> callback) {
+    LOG(INFO) << "Setting data receive callback...";
+    data_receive_callback_ = callback;
+    LOG(INFO) << "Data receive callback set successfully";
+}
+
+bool RpcCommunicator::initialize(const RpcCommunicatorConfig& config) {
+    config_ = config;
+    easylog::set_min_severity(easylog::Severity::WARNING);  // Set log level
+    // to WARNING
+
+    // Initialize client pools with proper configuration
+    coro_io::client_pool<coro_rpc::coro_rpc_client>::pool_config pool_conf{};
+    const char* value = std::getenv("MC_RPC_PROTOCOL");
+    if (value && std::string_view(value) == "rdma") {
+        pool_conf.client_config.socket_config =
+            coro_io::ib_socket_t::config_t{};
+    }
+    client_pools_ =
+        std::make_shared<coro_io::client_pools<coro_rpc::coro_rpc_client>>(
+            pool_conf);
+
+    LOG(INFO) << "create coro_rpc_client_pool with " << config.pool_size
+              << " threads";
+    if (!config.listen_address.empty()) {
+        LOG(INFO) << "Initializing server on " << config.listen_address;
+
+        server_ = std::make_unique<coro_rpc::coro_rpc_server>(
+            config.thread_count, config.listen_address,
+            std::chrono::seconds(config.timeout_seconds));
+
+        if (value && std::string_view(value) == "rdma") {
+            if (server_) {
+                try {
+                    server_->init_ibv();
+                    LOG(INFO) << "RDMA initialized successfully";
+                } catch (const std::exception& e) {
+                    LOG(ERROR) << "RDMA initialization failed: " << e.what();
+                    LOG(WARNING) << "Falling back to TCP mode";
+                    // Continue without RDMA - the server will use TCP
+                } catch (...) {
+                    LOG(ERROR)
+                        << "RDMA initialization failed with unknown error";
+                    LOG(WARNING) << "Falling back to TCP mode";
+                    // Continue without RDMA - the server will use TCP
+                }
+            } else {
+                LOG(ERROR) << "Server pointer is null, cannot initialize RDMA";
+                LOG(WARNING) << "Falling back to TCP mode";
+            }
+        }
+
+        server_->register_handler<&RpcCommunicator::handleDataTransfer,
+                                  &RpcCommunicator::handleTensorTransfer>(this);
+    }
+    LOG(INFO) << "Environment variable MC_RPC_PROTOCOL is set to "
+              << (value ? value : "not set");
+    if (value && std::string_view(value) == "rdma") {
+        LOG(INFO) << "Using RDMA transport for RPC communication";
+    } else {
+        LOG(INFO) << "Using TCP transport for RPC communication";
+    }
+
+    LOG(INFO) << "Communicator initialized with client pool support";
+    return true;
+}
+
+bool RpcCommunicator::startServerImpl(bool is_async) {
+    if (is_async) {
+        return this->startServerAsync();
+    } else {
+        return this->startServer();
+    }
+}
+
+bool RpcCommunicator::startServer() {
+    if (!server_ || config_.listen_address.empty()) return false;
+
+    try {
+        auto ec = server_->start();
+        if (ec.val() == 0) {
+            is_server_started_ = true;
+            LOG(INFO) << "Server started on " << config_.listen_address;
+            return true;
+        } else {
+            LOG(ERROR) << "Failed to start server: " << ec.message();
+            return false;
+        }
+    } catch (const std::exception& e) {
+        LOG(ERROR) << "Failed to start server: " << e.what();
+        return false;
+    }
+}
+
+bool RpcCommunicator::startServerAsync() {
+    if (!server_ || config_.listen_address.empty()) return false;
+
+    try {
+        auto ec = server_->async_start();
+        if (!ec.hasResult()) {
+            is_server_started_ = true;
+            LOG(INFO) << "Server started asynchronously on "
+                      << config_.listen_address;
+            return true;
+        } else {
+            LOG(ERROR) << "Failed to start server asynchronously";
+            return false;
+        }
+    } catch (const std::exception& e) {
+        LOG(ERROR) << "Failed to start server asynchronously: " << e.what();
+        return false;
+    }
+}
+
+void RpcCommunicator::stopServer() {
+    if (is_server_started_ && server_) {
+        server_->stop();
+        is_server_started_ = false;
+        LOG(INFO) << "Server stopped";
+    }
+}
+
+int RpcCommunicator::sendData(const std::string& target_address,
+                              const void* data, size_t data_size) {
+    auto result = async_simple::coro::syncAwait(
+        sendDataAsync(target_address, data, data_size));
+    return result.code;
+}
+
+async_simple::coro::Lazy<RpcResult> RpcCommunicator::sendDataAsync(
+    const std::string& target_address, const void* data, size_t data_size) {
+    std::string_view data_view(static_cast<const char*>(data), data_size);
+
+    // For large data, use attachment to avoid copying
+    const size_t ATTACHMENT_THRESHOLD = 1024;  // Use attachment for data > 1KB
+
+    auto rpc_result = co_await client_pools_->send_request(
+        target_address,
+        [data_view, data_size](coro_rpc::coro_rpc_client& client)
+            -> async_simple::coro::Lazy<void> {
+            if (data_size > ATTACHMENT_THRESHOLD) {
+                // Use attachment for large data - zero copy
+                client.set_req_attachment(data_view);
+                // Send empty data parameter, actual data in attachment
+                auto result =
+                    co_await client.call<&RpcCommunicator::handleDataTransfer>(
+                        std::string_view{});
+                if (!result.has_value()) {
+                    LOG(ERROR) << "RPC call failed: " << result.error().msg;
+                }
+            } else {
+                // Use regular parameter for small data
+                auto result =
+                    co_await client.call<&RpcCommunicator::handleDataTransfer>(
+                        data_view);
+                if (!result.has_value()) {
+                    LOG(ERROR) << "RPC call failed: " << result.error().msg;
+                }
+            }
+        });
+
+    if (!rpc_result.has_value()) {
+        LOG(ERROR) << "RPC send request failed";
+        co_return RpcResult{-1, "RPC call failed"};
+    }
+    RpcResult res;
+    res.code = 0;
+    co_return res;
+}
+
+int RpcCommunicator::sendTensor(const std::string& target_address,
+                                const pybind11::object& tensor) {
+    try {
+        TensorInfo tensor_info;
+        {
+            pybind11::gil_scoped_acquire acquire;
+            pybind11::object tensor_obj =
+                pybind11::reinterpret_borrow<pybind11::object>(tensor);
+
+            // Validate tensor type using duck typing - check for required
+            // attributes
+            if (!pybind11::hasattr(tensor_obj, "data_ptr") ||
+                !pybind11::hasattr(tensor_obj, "numel") ||
+                !pybind11::hasattr(tensor_obj, "element_size") ||
+                !pybind11::hasattr(tensor_obj, "shape") ||
+                !pybind11::hasattr(tensor_obj, "dtype")) {
+                LOG(ERROR) << "Input is not a valid tensor object (missing "
+                              "required attributes)";
+                return -1;
+            }
+
+            // Extract tensor properties
+            uintptr_t data_ptr =
+                tensor_obj.attr("data_ptr")().cast<uintptr_t>();
+            size_t numel = tensor_obj.attr("numel")().cast<size_t>();
+            size_t element_size =
+                tensor_obj.attr("element_size")().cast<size_t>();
+            size_t tensor_size = numel * element_size;
+
+            // Get tensor shape
+            pybind11::object shape_obj = tensor_obj.attr("shape");
+            pybind11::tuple shape_tuple =
+                pybind11::cast<pybind11::tuple>(shape_obj);
+            const size_t shape_size = shape_tuple.size();
+            std::vector<size_t> shape;
+            shape.reserve(shape_size);  // Pre-allocate to avoid reallocations
+            for (size_t i = 0; i < shape_size; ++i) {
+                shape.push_back(shape_tuple[i].cast<size_t>());
+            }
+
+            // Get tensor dtype string
+            pybind11::object dtype_obj = tensor_obj.attr("dtype");
+            std::string dtype = pybind11::str(dtype_obj).cast<std::string>();
+
+            // Fill TensorInfo structure (no data copying, just metadata)
+            tensor_info.data_ptr = reinterpret_cast<void*>(data_ptr);
+            tensor_info.total_bytes = tensor_size;
+            tensor_info.shape = std::move(shape);
+            tensor_info.dtype = std::move(dtype);
+
+            // Format shape as string for logging - optimize string building
+            // with reserve (glog LOG macro handles log level checking)
+            std::string shape_str;
+            // Estimate size: each number ~10 chars + ", " = ~12 chars per
+            // dimension
+            shape_str.reserve(tensor_info.shape.size() * 12);
+            shape_str = "[";
+            for (size_t i = 0; i < tensor_info.shape.size(); ++i) {
+                if (i > 0) shape_str += ", ";
+                shape_str += std::to_string(tensor_info.shape[i]);
+            }
+            shape_str += "]";
+            LOG(INFO) << "Sending tensor with shape: " << shape_str
+                      << ", dtype: " << tensor_info.dtype
+                      << ", tensor size: " << tensor_info.total_bytes
+                      << " bytes";
+        }
+
+        // Use the async version which supports zero-copy via attachments
+        pybind11::gil_scoped_release release;
+        auto result = async_simple::coro::syncAwait(
+            sendTensorAsync(target_address, tensor_info));
+        return result;
+
+    } catch (const std::exception& e) {
+        LOG(ERROR) << "Send tensor error: " << e.what();
+        return -1;
+    }
+}
+
+async_simple::coro::Lazy<int> RpcCommunicator::sendTensorAsync(
+    const std::string& target_address, const TensorInfo& tensor) {
+    auto rpc_result = co_await client_pools_->send_request(
+        target_address,
+        [tensor](coro_rpc::coro_rpc_client& client)
+            -> async_simple::coro::Lazy<void> {
+            client.set_req_attachment(std::string_view(
+                static_cast<const char*>(tensor.data_ptr), tensor.total_bytes));
+
+            auto result =
+                co_await client.call<&RpcCommunicator::handleTensorTransfer>();
+
+            if (!result.has_value()) {
+                LOG(ERROR) << "Tensor RPC call failed: " << result.error().msg;
+            }
+        });
+    if (!rpc_result.has_value()) {
+        LOG(ERROR) << "Tensor RPC send request failed";
+        co_return -1;
+    }
+    co_return 0;
+}
+
+int RpcCommunicator::receiveData(const std::string& source_address,
+                                 void* buffer, size_t buffer_size,
+                                 int timeout_ms) {
+    auto result = async_simple::coro::syncAwait(
+        receiveDataAsync(source_address, timeout_ms));
+    return 0;
+}
+
+async_simple::coro::Lazy<std::string> RpcCommunicator::receiveDataAsync(
+    const std::string& source_address, int timeout_ms) {
+    // For attachment-based data reception, we should use a different approach
+    // This method is typically called from the handler when data is received
+    // The actual data reception is handled by the registered handlers
+    co_return std::string();
+}  // Data reception is handled via context and attachment in handlers
+
+void RpcCommunicator::handleDataTransfer(coro_rpc::context<void> context,
+                                         std::string_view data) {
+    // Check if there's an attachment for large data
+    auto ctx_info = context.get_context_info();
+    auto attachment = ctx_info->get_request_attachment();
+
+    LOG(INFO) << "Handling data transfer - Data: " << data.size()
+              << " bytes, Attachment: " << attachment.size() << " bytes";
+    // Call the data receive callback if set
+    if (data_receive_callback_) {
+        LOG(INFO) << "Calling data receive callback...";
+        // Note: coro_rpc context doesn't provide get_remote_endpoint()
+        // Using empty string as placeholder - can be enhanced if needed
+        std::string_view source_address = "";
+
+        // Use attachment if available (for large data), otherwise use data
+        // parameter
+        if (!attachment.empty()) {
+            // Use attachment data directly without copying - zero copy approach
+            std::string_view attachment_view = attachment;
+            data_receive_callback_(source_address, attachment_view);
+        } else {
+            // For small data, use the regular data parameter
+            data_receive_callback_(source_address, data);
+        }
+    } else {
+        LOG(INFO) << "No data receive callback set!";
+    }
+
+    // Echo back the attachment for response (zero-copy)
+    if (!attachment.empty()) {
+        ctx_info->set_response_attachment(std::string_view("ok"));
+    }
+
+    context.response_msg();
+}
+
+void RpcCommunicator::handleTensorTransfer(coro_rpc::context<void> context) {
+    auto ctx_info = context.get_context_info();
+    auto attachment = ctx_info->get_request_attachment();
+
+    LOG(INFO) << "Handling tensor transfer: " << attachment.size() << " bytes";
+
+    // Call the data receive callback if set (tensor data is received via
+    // attachment)
+    if (data_receive_callback_) {
+        LOG(INFO) << "Calling data receive callback for tensor...";
+        // Note: coro_rpc context doesn't provide get_remote_endpoint()
+        // Using empty string as placeholder - can be enhanced if needed
+        std::string_view source_address = "";
+
+        // Pass the attachment data to the callback
+        data_receive_callback_(source_address, attachment);
+    } else {
+        LOG(INFO) << "No data receive callback set for tensor!";
+    }
+
+    ctx_info->set_response_attachment(attachment);
+    context.response_msg();
+}
+
+void RpcCommunicator::handleDataTransferWithAttachment(
+    coro_rpc::context<void> context, std::string_view data) {
+    py_rpc_context t{};
+    t.context_ = std::move(context);
+    py::gil_scoped_acquire acquire;
+    auto view =
+        py::memoryview::from_buffer(data.data(), {data.size()}, {sizeof(char)});
+
+    py_callback_(std::move(t), view);
+}
+
+void RpcCommunicator::handleTensorTransferWithAttachment(
+    coro_rpc::context<void> context) {
+    py_rpc_context t{};
+
+    // Get the attachment before moving the context
+    auto ctx_info = context.get_context_info();
+    auto attachment = ctx_info->get_request_attachment();
+
+    t.context_ = std::move(context);
+    py::gil_scoped_acquire acquire;
+
+    auto view = py::memoryview::from_buffer(
+        attachment.data(), {attachment.size()}, {sizeof(int8_t)});
+
+    py_callback_(std::move(t), view);
+}
+
+}  // namespace mooncake

mooncake-transfer-engine/src/transport/rpc_communicator/rpc_interface.cpp

@@ -0,0 +1,589 @@
+#include "transport/rpc_communicator/rpc_interface.h"
+#include "transport/rpc_communicator/rpc_communicator.h"
+#include "config.h"
+#include <memory>
+#include <thread>
+#include <future>
+#include <vector>
+#include <glog/logging.h>
+#include "async_simple/coro/SyncAwait.h"
+
+namespace mooncake {
+
+static constexpr size_t MAX_TENSOR_DIMS = 4;
+// Tensor metadata size: dtype (4 bytes) + ndim (4 bytes) + shape dimensions
+// (MAX_TENSOR_DIMS * 8 bytes)
+static constexpr size_t TENSOR_METADATA_SIZE = 4 + 4 + MAX_TENSOR_DIMS * 8;
+
+// Implementation class
+class RpcInterface::Impl {
+   public:
+    std::unique_ptr<RpcCommunicator> communicator;
+    pybind11::function data_receive_callback;
+    pybind11::function tensor_receive_callback;
+};
+
+// Constructor
+RpcInterface::RpcInterface() : impl_(std::make_unique<Impl>()) {}
+
+// Destructor
+RpcInterface::~RpcInterface() = default;
+
+// Initialize
+bool RpcInterface::initialize(const std::string& local_address,
+                              size_t thread_count, size_t timeout_seconds,
+                              size_t pool_size) {
+    RpcCommunicatorConfig config;
+    config.listen_address = local_address;
+    config.thread_count = thread_count;
+    config.timeout_seconds = timeout_seconds;
+    config.pool_size = pool_size;
+
+    impl_->communicator = std::make_unique<RpcCommunicator>();
+    return impl_->communicator->initialize(config);
+}
+
+// Convenience method for client initialization
+bool RpcInterface::initializeClient(size_t pool_size, size_t timeout_seconds) {
+    return initialize("", 0, timeout_seconds, pool_size);
+}
+
+// Convenience method for server initialization
+bool RpcInterface::initializeServer(const std::string& listen_address,
+                                    size_t thread_count,
+                                    size_t timeout_seconds) {
+    return initialize(listen_address, thread_count, timeout_seconds, 4);
+}
+
+bool RpcInterface::startServer() {
+    if (!impl_->communicator) return false;
+    return impl_->communicator->startServer();
+}
+
+bool RpcInterface::startServerAsync() {
+    if (!impl_->communicator) return false;
+    return impl_->communicator->startServerAsync();
+}
+
+bool RpcInterface::startServerImpl(bool is_async) {
+    if (!impl_->communicator) return false;
+    return impl_->communicator->startServerImpl(is_async);
+}
+
+void RpcInterface::stopServer() {
+    if (impl_->communicator) {
+        impl_->communicator->stopServer();
+    }
+}
+
+int RpcInterface::sendData(const std::string& target_address,
+                           pybind11::handle data) {
+    if (!impl_->communicator) return -1;
+
+    pybind11::gil_scoped_acquire acquire;
+
+    // Extract data from handle directly
+    std::string data_str;
+    try {
+        // Try to get direct string view from bytes object
+        pybind11::bytes data_bytes =
+            pybind11::reinterpret_borrow<pybind11::bytes>(data);
+        data_str = static_cast<std::string>(data_bytes);
+    } catch (...) {
+        // Fallback: convert to bytes and then get string
+        pybind11::bytes data_bytes = pybind11::cast<pybind11::bytes>(data);
+        data_str = static_cast<std::string>(data_bytes);
+    }
+
+    pybind11::gil_scoped_release release;
+    return impl_->communicator->sendData(target_address, data_str.data(),
+                                         data_str.size());
+}
+
+pybind11::object RpcInterface::sendDataAsync(const std::string& target_address,
+                                             pybind11::handle data,
+                                             pybind11::handle loop) {
+    pybind11::gil_scoped_acquire acquire;
+
+    auto future_module = pybind11::module_::import("asyncio");
+    auto future_obj = future_module.attr("Future")();
+
+    if (!impl_->communicator) {
+        auto exc_type =
+            pybind11::module_::import("builtins").attr("RuntimeError");
+        auto exc = exc_type(pybind11::str("Communicator not initialized"));
+        future_obj.attr("set_exception")(exc);
+        return future_obj;
+    }
+
+    auto communicator = impl_->communicator.get();
+    std::string target_addr = target_address;
+
+    // Extract data from handle directly without creating intermediate bytes
+    // object
+    std::string_view data_view;
+    std::string data_str;  // Only used if we can't get direct view
+
+    try {
+        // Try to get direct string view from bytes object
+        pybind11::bytes data_bytes =
+            pybind11::reinterpret_borrow<pybind11::bytes>(data);
+        data_view = data_bytes;
+        // Store a copy for lambda capture since string_view might not be valid
+        // after GIL release
+        data_str = std::string(data_view);
+    } catch (...) {
+        // Fallback: convert to bytes and then to string
+        pybind11::bytes data_bytes = pybind11::cast<pybind11::bytes>(data);
+        data_str = data_bytes;
+    }
+
+    // Release GIL before starting coroutine
+    pybind11::gil_scoped_release release;
+
+    auto coro_lambda = [communicator, target_addr, data_str, future_obj,
+                        loop]() -> async_simple::coro::Lazy<void> {
+        try {
+            auto result_struct = co_await communicator->sendDataAsync(
+                target_addr, data_str.data(), data_str.size());
+            int result = result_struct.code;
+
+            auto call_soon_threadsafe = [future_obj, loop, result]() {
+                pybind11::gil_scoped_acquire acquire;
+                if (result >= 0) {
+                    future_obj.attr("set_result")(result);
+                } else {
+                    auto runtime_error = pybind11::module_::import("builtins")
+                                             .attr("RuntimeError");
+                    future_obj.attr("set_exception")(
+                        runtime_error(pybind11::str("Send data failed")));
+                }
+            };
+
+            auto callback = pybind11::cpp_function(call_soon_threadsafe);
+            loop.attr("call_soon_threadsafe")(callback);
+        } catch (const std::exception& e) {
+            auto call_soon_threadsafe = [future_obj, loop, e]() {
+                pybind11::gil_scoped_acquire acquire;
+                auto runtime_error =
+                    pybind11::module_::import("builtins").attr("RuntimeError");
+                future_obj.attr("set_exception")(runtime_error(pybind11::str(
+                    std::string("Send data error: ") + e.what())));
+            };
+
+            auto callback = pybind11::cpp_function(call_soon_threadsafe);
+            loop.attr("call_soon_threadsafe")(callback);
+        }
+    };
+
+    auto lazy = coro_lambda();
+    lazy.start([](auto&& result) {
+        if (result.hasError()) {
+            LOG(ERROR) << "Coroutine completed with error";
+        }
+    });
+
+    return future_obj;
+}
+
+int RpcInterface::sendTensor(const std::string& target_address,
+                             pybind11::handle tensor) {
+    if (!impl_->communicator) return -1;
+
+    try {
+        TensorInfo tensor_info;
+        {
+            pybind11::gil_scoped_acquire acquire;
+            pybind11::object tensor_obj =
+                pybind11::reinterpret_borrow<pybind11::object>(tensor);
+
+            // Validate tensor type using duck typing - check for required
+            // attributes
+            if (!pybind11::hasattr(tensor_obj, "data_ptr") ||
+                !pybind11::hasattr(tensor_obj, "numel") ||
+                !pybind11::hasattr(tensor_obj, "element_size") ||
+                !pybind11::hasattr(tensor_obj, "shape") ||
+                !pybind11::hasattr(tensor_obj, "dtype")) {
+                LOG(ERROR) << "Input is not a valid tensor object (missing "
+                              "required attributes)";
+                return -1;
+            }
+
+            // Extract tensor properties - zero copy, just get pointers and
+            // metadata
+            uintptr_t data_ptr =
+                tensor_obj.attr("data_ptr")().cast<uintptr_t>();
+            size_t numel = tensor_obj.attr("numel")().cast<size_t>();
+            size_t element_size =
+                tensor_obj.attr("element_size")().cast<size_t>();
+            size_t tensor_size = numel * element_size;
+
+            // Get tensor shape
+            pybind11::object shape_obj = tensor_obj.attr("shape");
+            pybind11::tuple shape_tuple =
+                pybind11::cast<pybind11::tuple>(shape_obj);
+            std::vector<size_t> shape;
+            for (size_t i = 0; i < shape_tuple.size(); i++) {
+                shape.push_back(shape_tuple[i].cast<size_t>());
+            }
+
+            // Get tensor dtype string
+            pybind11::object dtype_obj = tensor_obj.attr("dtype");
+            std::string dtype = dtype_obj.attr("__str__")().cast<std::string>();
+
+            // Fill TensorInfo structure (no data copying, just metadata)
+            tensor_info.data_ptr = reinterpret_cast<void*>(data_ptr);
+            tensor_info.total_bytes = tensor_size;
+            tensor_info.shape = std::move(shape);
+            tensor_info.dtype = std::move(dtype);
+
+            // Format shape as string for logging
+            std::string shape_str = "[";
+            for (size_t i = 0; i < tensor_info.shape.size(); i++) {
+                shape_str += std::to_string(tensor_info.shape[i]);
+                if (i < tensor_info.shape.size() - 1) shape_str += ", ";
+            }
+            shape_str += "]";
+            LOG(INFO) << "Sending tensor with shape: " << shape_str
+                      << ", dtype: " << tensor_info.dtype
+                      << ", tensor size: " << tensor_info.total_bytes
+                      << " bytes";
+        }
+
+        // Use the async version which supports zero-copy via attachments
+        pybind11::gil_scoped_release release;
+        auto result = async_simple::coro::syncAwait(
+            impl_->communicator->sendTensorAsync(target_address, tensor_info));
+        return result;
+
+    } catch (const std::exception& e) {
+        LOG(ERROR) << "Send tensor error: " << e.what();
+        return -1;
+    }
+}
+
+pybind11::object RpcInterface::sendTensorAsync(
+    const std::string& target_address, pybind11::handle tensor,
+    pybind11::handle loop) {
+    pybind11::gil_scoped_acquire acquire;
+
+    auto future_module = pybind11::module_::import("asyncio");
+    auto future_obj = future_module.attr("Future")();
+
+    if (!impl_->communicator) {
+        auto exc_type =
+            pybind11::module_::import("builtins").attr("RuntimeError");
+        auto exc = exc_type(pybind11::str("Communicator not initialized"));
+        future_obj.attr("set_exception")(exc);
+        return future_obj;
+    }
+
+    auto communicator = impl_->communicator.get();
+    std::string target_addr = target_address;
+
+    // Extract tensor info
+    pybind11::object tensor_obj =
+        pybind11::reinterpret_borrow<pybind11::object>(tensor);
+    uintptr_t data_ptr = tensor_obj.attr("data_ptr")().cast<uintptr_t>();
+    size_t numel = tensor_obj.attr("numel")().cast<size_t>();
+    size_t element_size = tensor_obj.attr("element_size")().cast<size_t>();
+    size_t tensor_size = numel * element_size;
+
+    // Get tensor shape and dtype
+    pybind11::object shape_obj = tensor_obj.attr("shape");
+    pybind11::tuple shape_tuple = pybind11::cast<pybind11::tuple>(shape_obj);
+    std::vector<size_t> shape;
+    for (size_t i = 0; i < shape_tuple.size(); i++) {
+        shape.push_back(shape_tuple[i].cast<size_t>());
+    }
+
+    pybind11::object dtype_obj = tensor_obj.attr("dtype");
+    std::string dtype = dtype_obj.attr("__str__")().cast<std::string>();
+
+    // Create TensorInfo on stack (no need for shared_ptr)
+    TensorInfo tensor_info;
+    tensor_info.data_ptr = reinterpret_cast<void*>(data_ptr);
+    tensor_info.total_bytes = tensor_size;
+    tensor_info.shape = std::move(shape);
+    tensor_info.dtype = std::move(dtype);
+
+    // Keep a reference to the tensor object to avoid it being garbage collected
+    pybind11::object py_tensor_obj = tensor_obj;
+
+    pybind11::gil_scoped_release release;
+
+    auto coro_lambda = [communicator, target_addr, tensor_info, future_obj,
+                        loop,
+                        py_tensor_obj]() -> async_simple::coro::Lazy<void> {
+        try {
+            auto result = co_await communicator->sendTensorAsync(target_addr,
+                                                                 tensor_info);
+
+            auto call_soon_threadsafe = [future_obj, loop, result]() {
+                pybind11::gil_scoped_acquire acquire;
+                if (result >= 0) {
+                    future_obj.attr("set_result")(result);
+                } else {
+                    auto runtime_error = pybind11::module_::import("builtins")
+                                             .attr("RuntimeError");
+                    future_obj.attr("set_exception")(
+                        runtime_error(pybind11::str("Send tensor failed")));
+                }
+            };
+
+            auto callback = pybind11::cpp_function(call_soon_threadsafe);
+            loop.attr("call_soon_threadsafe")(callback);
+        } catch (const std::exception& e) {
+            auto call_soon_threadsafe = [future_obj, loop, e]() {
+                pybind11::gil_scoped_acquire acquire;
+                auto runtime_error =
+                    pybind11::module_::import("builtins").attr("RuntimeError");
+                future_obj.attr("set_exception")(runtime_error(pybind11::str(
+                    std::string("Send tensor error: ") + e.what())));
+            };
+
+            auto callback = pybind11::cpp_function(call_soon_threadsafe);
+            loop.attr("call_soon_threadsafe")(callback);
+        }
+    };
+
+    // Start the coroutine in a detached manner (fire and forget)
+    auto lazy = coro_lambda();
+    lazy.start([](auto&& result) {
+        // This callback will be called when the coroutine completes
+        // We don't need to do anything here since the result is handled in
+        // the coroutine itself
+        if (result.hasError()) {
+            // Log error if needed
+            LOG(ERROR) << "Tensor coroutine completed with error";
+        }
+    });
+
+    return future_obj;
+}
+
+void RpcInterface::setDataReceiveCallback(pybind11::function callback) {
+    pybind11::gil_scoped_acquire acquire;
+    impl_->data_receive_callback = callback;
+    if (impl_->communicator) {
+        // Capture this pointer - safe because RpcInterface is managed by Python
+        // and will remain valid as long as the callback is set
+        auto interface_ptr = this;
+        impl_->communicator->setDataReceiveCallback(
+            [interface_ptr](std::string_view source, std::string_view data) {
+                interface_ptr->handleIncomingData(source, data);
+            });
+    }
+}
+
+void RpcInterface::setTensorReceiveCallback(pybind11::function callback) {
+    pybind11::gil_scoped_acquire acquire;
+    impl_->tensor_receive_callback = callback;
+
+    // Note: Tensor data is received through the regular data callback
+    // The handleIncomingData function will detect tensor data and route it
+    // to handleIncomingTensor automatically
+}
+
+void RpcInterface::handleIncomingData(std::string_view source,
+                                      std::string_view data) {
+    LOG(INFO) << "RpcInterface::handleIncomingData called with " << data.size()
+              << " bytes";
+
+    // C++ tensor rebuilding
+    if (data.size() >= TENSOR_METADATA_SIZE) {
+        // Read the first few bytes to check if it looks like tensor metadata
+        const uint32_t* header = reinterpret_cast<const uint32_t*>(data.data());
+        uint32_t dtype = header[0];
+        uint32_t ndim = header[1];
+
+        LOG(INFO) << "Checking tensor metadata: dtype=" << dtype
+                  << ", ndim=" << ndim;
+
+        // Basic validation: check if dtype and ndim are in reasonable ranges
+        if (dtype > 0 && dtype <= 9 && ndim <= 4) {
+            LOG(INFO)
+                << "Data recognized as tensor, calling handleIncomingTensor";
+
+            // This looks like tensor data, handle it as such
+            std::vector<size_t> shape;
+            const int64_t* shape_data =
+                reinterpret_cast<const int64_t*>(data.data() + 8);
+            if (data.size() < TENSOR_METADATA_SIZE ||
+                data.size() < 8 + ndim * sizeof(int64_t)) {
+                LOG(WARNING) << "Data too small for claimed tensor metadata";
+                return;
+            }
+            for (int i = 0; i < static_cast<int>(ndim); i++) {
+                if (shape_data[i] > 0 && shape_data[i] < (1LL << 48)) {
+                    shape.push_back(static_cast<size_t>(shape_data[i]));
+                } else if (shape_data[i] > 0) {
+                    LOG(WARNING) << "Shape dimension " << i
+                                 << " too large: " << shape_data[i];
+                    return;
+                }
+            }
+
+            // Get dtype name based on dtype ID
+            std::string_view dtype_name;
+            switch (dtype) {
+                case 1:
+                    dtype_name = "float16";
+                    break;
+                case 2:
+                    dtype_name = "float32";
+                    break;
+                case 3:
+                    dtype_name = "float64";
+                    break;
+                case 4:
+                    dtype_name = "int8";
+                    break;
+                case 5:
+                    dtype_name = "int16";
+                    break;
+                case 6:
+                    dtype_name = "int32";
+                    break;
+                case 7:
+                    dtype_name = "int64";
+                    break;
+                case 8:
+                    dtype_name = "uint8";
+                    break;
+                case 9:
+                    dtype_name = "bool";
+                    break;
+                default:
+                    dtype_name = "unknown";
+                    break;
+            }
+
+            // Call tensor handler instead of data handler
+            handleIncomingTensor(source, data, shape, dtype_name);
+            return;
+        }
+    }
+
+    // Handle as regular data if not tensor data
+    if (!impl_->data_receive_callback) return;
+
+    try {
+        pybind11::gil_scoped_acquire acquire;
+        pybind11::dict received;
+        received["source"] =
+            std::string(source);  // Convert to string for Python
+        received["data"] = pybind11::bytes(
+            std::string(data));  // Convert to string for pybind11::bytes
+
+        impl_->data_receive_callback(received);
+    } catch (const std::exception& e) {
+        LOG(ERROR) << "Error in data receive callback: " << e.what();
+    }
+}
+
+void RpcInterface::handleIncomingTensor(std::string_view source,
+                                        std::string_view data,
+                                        const std::vector<size_t>& shape,
+                                        std::string_view dtype) {
+    LOG(INFO) << "RpcInterface::handleIncomingTensor called"
+              << " - source: " << source << ", data size: " << data.size()
+              << ", dtype: " << dtype << ", shape size: " << shape.size();
+
+    if (!impl_->tensor_receive_callback) {
+        LOG(WARNING) << "No tensor receive callback set!";
+        return;
+    }
+
+    LOG(INFO) << "Calling Python tensor receive callback...";
+
+    try {
+        pybind11::gil_scoped_acquire acquire;
+
+        ReceivedTensor received;
+        received.source_address = std::string(source);
+        received.data = std::string(data);
+        received.shape = shape;
+        received.dtype = std::string(dtype);
+
+        impl_->tensor_receive_callback(received);
+    } catch (const std::exception& e) {
+        LOG(ERROR) << "Error in tensor receive callback: " << e.what();
+    }
+}
+
+// Factory functions for creating RPC client and server
+std::unique_ptr<RpcInterface> createRpcClient(uint64_t local_rank,
+                                              uint64_t world_size) {
+    auto client = std::make_unique<RpcInterface>();
+    // Initialize client with default settings
+    client->initialize("", 0, 30, 10);
+    return client;
+}
+
+std::unique_ptr<RpcInterface> createRpcServer(uint64_t local_rank,
+                                              uint64_t world_size) {
+    auto server = std::make_unique<RpcInterface>();
+    // Initialize server with default settings
+    server->initialize("0.0.0.0:8080", 0, 30, 10);
+    return server;
+}
+
+// Python binding implementation
+void bind_rpc_interface(pybind11::module_& m) {
+    namespace py = pybind11;
+    using namespace mooncake;
+
+    // Bind RpcInterface::ReceivedData
+    py::class_<RpcInterface::ReceivedData>(m, "ReceivedData")
+        .def_readonly("source_address",
+                      &RpcInterface::ReceivedData::source_address)
+        .def_readonly("data_size", &RpcInterface::ReceivedData::data_size)
+        .def("get_bytes", &RpcInterface::ReceivedData::getBytes)
+        .def("get_memory_view", &RpcInterface::ReceivedData::getMemoryView);
+
+    // Bind RpcInterface::ReceivedTensor
+    py::class_<RpcInterface::ReceivedTensor>(m, "ReceivedTensor")
+        .def_readonly("source_address",
+                      &RpcInterface::ReceivedTensor::source_address)
+        .def_readonly("shape", &RpcInterface::ReceivedTensor::shape)
+        .def_readonly("dtype", &RpcInterface::ReceivedTensor::dtype)
+        .def_readonly("total_bytes", &RpcInterface::ReceivedTensor::total_bytes)
+        .def("get_data_size", &RpcInterface::ReceivedTensor::getDataSize)
+        .def("get_data_as_bytes", &RpcInterface::ReceivedTensor::getDataAsBytes)
+        .def("get_memory_view", &RpcInterface::ReceivedTensor::getMemoryView);
+
+    // Bind RpcInterface
+    py::class_<RpcInterface>(m, "RpcInterface")
+        .def(py::init<>())
+        .def("initialize", &RpcInterface::initialize,
+             py::arg("listen_address") = "", py::arg("thread_count") = 0,
+             py::arg("timeout_seconds") = 30, py::arg("pool_size") = 10)
+        .def("initialize_client", &RpcInterface::initializeClient,
+             py::arg("pool_size") = 10, py::arg("timeout_seconds") = 30)
+        .def("initialize_server", &RpcInterface::initializeServer,
+             py::arg("listen_address"), py::arg("thread_count") = 8,
+             py::arg("timeout_seconds") = 30)
+        .def("start_server", &RpcInterface::startServer)
+        .def("start_server_async", &RpcInterface::startServerAsync)
+        .def("stop_server", &RpcInterface::stopServer)
+        .def("send_data", &RpcInterface::sendData, py::arg("target_address"),
+             py::arg("data"))
+        .def("send_data_async", &RpcInterface::sendDataAsync,
+             py::arg("target_address"), py::arg("data"), py::arg("loop"))
+        .def("send_tensor", &RpcInterface::sendTensor,
+             py::arg("target_address"), py::arg("tensor"))
+        .def("send_tensor_async", &RpcInterface::sendTensorAsync,
+             py::arg("target_address"), py::arg("tensor"), py::arg("loop"))
+        .def("set_data_receive_callback", &RpcInterface::setDataReceiveCallback)
+        .def("set_tensor_receive_callback",
+             &RpcInterface::setTensorReceiveCallback);
+
+    // Bind factory functions
+    m.def("create_rpc_client", &createRpcClient, py::arg("local_rank") = 0,
+          py::arg("world_size") = 1);
+    m.def("create_rpc_server", &createRpcServer, py::arg("local_rank") = 0,
+          py::arg("world_size") = 1);
+}
+
+}  // namespace mooncake

mooncake-transfer-engine/tests/rpc_communicator_test.py

@@ -0,0 +1,134 @@
+import time
+import threading
+import argparse
+import mooncake.engine as engine
+
+class AtomicCounter:
+    def __init__(self, initial=0):
+        self._value = initial
+        self._lock = threading.Lock()
+
+    def inc(self, num=1):
+        with self._lock:
+            self._value += num
+            return self._value
+
+    def dec(self, num=1):
+        with self._lock:
+            self._value -= num
+            return self._value
+
+    def get_and_reset(self):
+        with self._lock:
+            r = self._value
+            self._value = 0
+            return r
+
+counter = AtomicCounter()
+
+# Global variable to store data size
+data_size = 1024 * 1024  # Default 1MB
+test_data = None
+
+def print_qps():
+    while True:
+        time.sleep(1)
+        val = counter.get_and_reset()
+        if val == 0:
+            continue
+        print("bandwidth:", 8 * val * data_size / (1024*1024*1024), "GB/s")
+
+def send_data(client, target_url):
+    while True:
+        try:
+            result = client.send_data(target_url, test_data)
+            if result < 0:
+                print(f"Warning: send_data returned error code {result}")
+                time.sleep(0.01)  # Brief delay on error to avoid tight loop
+            else:
+                counter.inc()
+        except Exception as e:
+            print(f"Error sending data to {target_url}: {e}")
+            time.sleep(0.1)  # Delay on exception to avoid rapid retry
+
+def run_server(bind_url, data_size_mb=1):
+    """Run server mode"""
+    global data_size, test_data
+    data_size = data_size_mb * 1024 * 1024
+    test_data = b'\x00' * data_size
+    
+    print(f"Starting server on {bind_url} with {data_size_mb}MB data packets")
+    
+    RpcInterface = engine.RpcInterface
+    server = RpcInterface()
+    server.initialize_server(bind_url, thread_count=8)
+    server.start_server_async()
+    
+    # Start QPS statistics thread
+    thread = threading.Thread(target=print_qps)
+    thread.daemon = True
+    thread.start()
+    
+    print(f"Server started on {bind_url}, press Ctrl+C to stop")
+    try:
+        while True:
+            time.sleep(1)
+    except KeyboardInterrupt:
+        print("\nServer stopping...")
+        server.stop_server()
+
+def run_client(target_url, num_threads=8, data_size_mb=1):
+    """Run client mode"""
+    global data_size, test_data
+    data_size = data_size_mb * 1024 * 1024
+    test_data = b'\x00' * data_size
+    
+    print(f"Starting client, connecting to {target_url} with {num_threads} threads, {data_size_mb}MB data packets")
+    
+    RpcInterface = engine.RpcInterface
+    client = RpcInterface()
+    client.initialize_client(pool_size=100)
+    
+    # Start QPS statistics thread
+    qps_thread = threading.Thread(target=print_qps)
+    qps_thread.daemon = True
+    qps_thread.start()
+    
+    # Start multiple sending threads
+    threads = []
+    for i in range(num_threads):
+        thread = threading.Thread(target=send_data, args=(client, target_url))
+        thread.daemon = True
+        thread.start()
+        threads.append(thread)
+    
+    print(f"Client started with {num_threads} threads, press Ctrl+C to stop")
+    try:
+        while True:
+            time.sleep(1)
+    except KeyboardInterrupt:
+        print("\nClient stopping...")
+
+def main():
+    parser = argparse.ArgumentParser(description='Mooncake Communication Bandwidth Test Tool')
+    parser.add_argument('mode', choices=['server', 'client'], 
+                       help='Run mode: server or client')
+    parser.add_argument('--url', default='127.0.0.1:9004',
+                       help='URL address (default: 127.0.0.1:9004)')
+    parser.add_argument('--threads', type=int, default=8,
+                       help='Number of client threads (default: 8)')
+    parser.add_argument('--data-size', type=int, default=1,
+                       help='Data packet size in MB (default: 1)')
+
+    args = parser.parse_args()
+    
+    if args.mode == 'server':
+        # Server mode, URL as bind address
+        bind_url = f"0.0.0.0:{args.url.split(':')[-1]}" if ':' in args.url else f"0.0.0.0:{args.url}"
+        run_server(bind_url, args.data_size)
+    else:
+        # Client mode, URL as target address
+        run_client(args.url, args.threads, args.data_size)
+
+if __name__ == "__main__":
+    main()