mooncake-track
/
Mooncake
0
0
Fork 12
Code Issues Pull Requests 13 Actions Packages Projects Releases Wiki Activity

feat(TE): Add MC_ENABLE_PARALLEL_REG_MR Option (#1238)

This commit is contained in:
Stary 2025-12-20 19:19:24 +08:00 committed by GitHub
parent cae1728e9f
commit 87c14b8c97
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 141 additions and 25 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
docs/source/design/transfer-engine/index.md
View File

@ -442,6 +442,7 @@ For advanced users, TransferEngine provides the following advanced runtime optio
- `MC_REDIS_DB_INDEX` The database index for Redis storage plugin, must be an integer between 0 and 255. Only takes effect when Redis is specified as the metadata server. If not set or invalid, the default value is 0.
- `MC_FRAGMENT_RATIO ` In RdmaTransport::submitTransferTask, if the last data piece after division is ≤ 1/MC_FRAGMENT_RATIO of the block size, it merges with the previous block to reduce overhead. The default value is 4
- `MC_ENABLE_DEST_DEVICE_AFFINITY` Enable device affinity for RDMA performance optimization. When enabled, Transfer Engine will prioritize communication with remote NICs that have the same name as local NICs to reduce QP count and improve network performance in rail-optimized topologies. The default value is false
- `MC_ENABLE_PARALLEL_REG_MR` Control parallel memory region registration across multiple RDMA NICs. Valid values: -1 (auto, default), 0 (disabled), 1 (enabled). When set to -1, parallel registration is automatically enabled when multiple RNICs exist and memory has been pre-touched. Note: If memory hasn't been touched before registration, parallel registration can be slower than sequential registration
- `MC_FORCE_MNNVL` Force to use Multi-Node NVLink as the active transport regardless whether RDMA devices are installed.
- `MC_FORCE_TCP` Force to use TCP as the active transport regardless whether RDMA devices are installed.
- `MC_MIN_PRC_PORT` Specifies the minimum port number for RPC service. The default value is 15000.

1
mooncake-transfer-engine/include/config.h
View File

@ -56,6 +56,7 @@ struct GlobalConfig {
bool use_ipv6 = false;
size_t fragment_limit = 16384;
bool enable_dest_device_affinity = false;
int parallel_reg_mr = -1;
size_t eic_max_block_size = 64UL * 1024 * 1024;
EndpointStoreType endpoint_store_type = EndpointStoreType::SIEVE;
int ib_traffic_class = -1;

11
mooncake-transfer-engine/include/transport/rdma_transport/rdma_transport.h
View File

@ -71,6 +71,17 @@ class RdmaTransport : public Transport {
int unregisterLocalMemoryBatch(
const std::vector<void *> &addr_list) override;
private:
// Internal version with force_sequential option to avoid nested parallelism
int registerLocalMemoryInternal(void *addr, size_t length,
const std::string &location,
bool remote_accessible,
bool update_metadata,
bool force_sequential);
int unregisterLocalMemoryInternal(void *addr, bool update_metadata,
bool force_sequential);
// TRANSFER
Status submitTransfer(BatchID batch_id,

13
mooncake-transfer-engine/src/config.cpp
View File

@ -293,6 +293,18 @@ void loadGlobalConfig(GlobalConfig &config) {
config.enable_dest_device_affinity = true;
}
const char *enable_parallel_reg_mr =
std::getenv("MC_ENABLE_PARALLEL_REG_MR");
if (enable_parallel_reg_mr) {
int val = atoi(enable_parallel_reg_mr);
if (val >= -1 && val <= 1) {
config.parallel_reg_mr = val;
} else {
LOG(WARNING) << "Ignore value from environment variable "
"MC_ENABLE_PARALLEL_REG_MR";
}
}
const char *endpoint_store_type_env = std::getenv("MC_ENDPOINT_STORE_TYPE");
if (endpoint_store_type_env) {
if (strcmp(endpoint_store_type_env, "FIFO") == 0) {
@ -381,6 +393,7 @@ void dumpGlobalConfig() {
LOG(INFO) << "max_wr = " << config.max_wr;
LOG(INFO) << "max_inline = " << config.max_inline;
LOG(INFO) << "mtu_length = " << mtuLengthToString(config.mtu_length);
LOG(INFO) << "parallel_reg_mr = " << config.parallel_reg_mr;
LOG(INFO) << "ib_traffic_class = " << config.ib_traffic_class;
}

140
mooncake-transfer-engine/src/transport/rdma_transport/rdma_transport.cpp
View File

@ -19,6 +19,7 @@
#include <sys/time.h>
#include <cassert>
#include <chrono>
#include <cstddef>
#include <future>
#include <set>
@ -174,6 +175,15 @@ int RdmaTransport::registerLocalMemory(void *addr, size_t length,
const std::string &name,
bool remote_accessible,
bool update_metadata) {
return registerLocalMemoryInternal(addr, length, name, remote_accessible,
update_metadata, false);
}
int RdmaTransport::registerLocalMemoryInternal(void *addr, size_t length,
const std::string &name,
bool remote_accessible,
bool update_metadata,
bool force_sequential) {
(void)remote_accessible;
BufferDesc buffer_desc;
const int kBaseAccessRights = IBV_ACCESS_LOCAL_WRITE |
@ -195,19 +205,35 @@ int RdmaTransport::registerLocalMemory(void *addr, size_t length,
}
}
if (context_list_.size() > 1 && do_pre_touch) {
// Parallel register the memory region. If the memory address has not
// been touched before, parallel register will be much slower than
// sequential register.
// Details in: https://github.com/kvcache-ai/Mooncake/issues/848
/* Parallel register when:
1. parallel_reg_mr is enabled via MC_ENABLE_PARALLEL_REG_MR;
2. parallel_reg_mr not set and multiple contexts exist and memory has been
pre-touched
Note: If memory hasn't been touched, parallel register can be
slower. Details in: https://github.com/kvcache-ai/Mooncake/issues/848
Note: force_sequential is used by batch operations to avoid nested
parallelism.
*/
int use_parallel_reg = 0;
if (!force_sequential) {
use_parallel_reg = globalConfig().parallel_reg_mr;
if (use_parallel_reg == -1) {
use_parallel_reg = context_list_.size() > 1 && do_pre_touch;
}
}
auto reg_start = std::chrono::steady_clock::now();
if (use_parallel_reg) {
std::vector<std::thread> reg_threads;
reg_threads.reserve(context_list_.size());
std::vector<int> ret_codes(context_list_.size(), 0);
const int ar = access_rights; // Local copy for lambda capture
for (size_t i = 0; i < context_list_.size(); ++i) {
reg_threads.emplace_back([this, &ret_codes, i, addr, length]() {
ret_codes[i] = context_list_[i]->registerMemoryRegion(
addr, length, access_rights);
reg_threads.emplace_back([this, &ret_codes, i, addr, length, ar]() {
ret_codes[i] =
context_list_[i]->registerMemoryRegion(addr, length, ar);
});
}
@ -217,17 +243,37 @@ int RdmaTransport::registerLocalMemory(void *addr, size_t length,
for (size_t i = 0; i < ret_codes.size(); ++i) {
if (ret_codes[i] != 0) {
LOG(ERROR) << "Failed to register memory region with context "
<< i;
return ret_codes[i];
}
}
} else {
for (auto &context : context_list_) {
int ret =
context->registerMemoryRegion(addr, length, access_rights);
if (ret) return ret;
for (size_t i = 0; i < context_list_.size(); ++i) {
int ret = context_list_[i]->registerMemoryRegion(addr, length,
access_rights);
if (ret) {
LOG(ERROR) << "Failed to register memory region with context "
<< i;
return ret;
}
}
}
auto reg_end = std::chrono::steady_clock::now();
auto reg_duration_ms =
std::chrono::duration_cast<std::chrono::milliseconds>(reg_end -
reg_start)
.count();
if (globalConfig().trace) {
LOG(INFO) << "registerMemoryRegion: addr=" << addr
<< ", length=" << length
<< ", contexts=" << context_list_.size()
<< ", parallel=" << (use_parallel_reg ? "true" : "false")
<< ", duration=" << reg_duration_ms << "ms";
}
// Collect keys from all contexts
for (auto &context : context_list_) {
buffer_desc.lkey.push_back(context->lkey(addr));
@ -242,26 +288,68 @@ int RdmaTransport::registerLocalMemory(void *addr, size_t length,
getMemoryLocation(addr, length, only_first_page);
if (entries.empty()) return -1;
buffer_desc.name = entries[0].location;
buffer_desc.addr = (uint64_t)addr;
buffer_desc.length = length;
int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata);
if (rc) return rc;
} else {
buffer_desc.name = name;
buffer_desc.addr = (uint64_t)addr;
buffer_desc.length = length;
int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata);
if (rc) return rc;
}
buffer_desc.addr = (uint64_t)addr;
buffer_desc.length = length;
int rc = metadata_->addLocalMemoryBuffer(buffer_desc, update_metadata);
if (rc) return rc;
return 0;
}
int RdmaTransport::unregisterLocalMemory(void *addr, bool update_metadata) {
return unregisterLocalMemoryInternal(addr, update_metadata, false);
}
int RdmaTransport::unregisterLocalMemoryInternal(void *addr,
bool update_metadata,
bool force_sequential) {
int rc = metadata_->removeLocalMemoryBuffer(addr, update_metadata);
if (rc) return rc;
for (auto &context : context_list_) context->unregisterMemoryRegion(addr);
// force_sequential is used by batch operations to avoid nested parallelism
int use_parallel_unreg = 0;
if (!force_sequential) {
use_parallel_unreg = globalConfig().parallel_reg_mr;
if (use_parallel_unreg == -1) {
use_parallel_unreg = context_list_.size() > 1;
}
}
if (use_parallel_unreg) {
std::vector<std::thread> unreg_threads;
unreg_threads.reserve(context_list_.size());
std::vector<int> ret_codes(context_list_.size(), 0);
for (size_t i = 0; i < context_list_.size(); ++i) {
unreg_threads.emplace_back([this, &ret_codes, i, addr]() {
ret_codes[i] = context_list_[i]->unregisterMemoryRegion(addr);
});
}
for (auto &thread : unreg_threads) {
thread.join();
}
for (size_t i = 0; i < ret_codes.size(); ++i) {
if (ret_codes[i] != 0) {
LOG(ERROR) << "Failed to unregister memory region with context "
<< i;
return ret_codes[i];
}
}
} else {
for (size_t i = 0; i < context_list_.size(); ++i) {
int ret = context_list_[i]->unregisterMemoryRegion(addr);
if (ret) {
LOG(ERROR) << "Failed to unregister memory region with context "
<< i;
return ret;
}
}
}
return 0;
}
@ -291,8 +379,9 @@ int RdmaTransport::registerLocalMemoryBatch(
for (auto &buffer : buffer_list) {
results.emplace_back(
std::async(std::launch::async, [this, buffer, location]() -> int {
return registerLocalMemory(buffer.addr, buffer.length, location,
true, false);
// Use force_sequential=true to avoid nested parallelism
return registerLocalMemoryInternal(buffer.addr, buffer.length,
location, true, false, true);
}));
}
@ -313,7 +402,8 @@ int RdmaTransport::unregisterLocalMemoryBatch(
for (auto &addr : addr_list) {
results.emplace_back(
std::async(std::launch::async, [this, addr]() -> int {
return unregisterLocalMemory(addr, false);
// Use force_sequential=true to avoid nested parallelism
return unregisterLocalMemoryInternal(addr, false, true);
}));
}