kvcache-ai / kvcache-ai/Mooncake

[Feature Request]: Topology- and load-aware remote replica selection in SelectBestReplica

Open
#2,516 8 comments 0 reactions 1 assignee Claimed by @catyans View on GitHub
Dominant language
C++
Stars
6.6k
Forks
1.2k
Avg merge
3d 5h
Merged PRs (30d)
312

Description

### Describe your feature request

**Problem**

When the master returns multiple remote MEMORY replicas for the same key, `SelectBestReplica` in `mooncake-store/src/real_client.cpp` (lines 290-326) picks the **first** one it encounters rather than the closest/least-loaded one:

```cpp
// real_client.cpp:293-314
const Replica::Descriptor *first_memory = nullptr;
const Replica::Descriptor *first_nof = nullptr;
for (const auto &r : replicas) {
if (r.status != ReplicaStatus::COMPLETE) continue;
if (r.is_memory_replica()) {
if (local_endpoints.count(...)) return &r; // local MEMORY — best case
if (!first_memory) first_memory = &r; // <-- blindly keeps the FIRST remote one
}
...
}
if (first_memory) return first_memory; // <-- returned without comparing peers
```

The current selection logic is purely **type + locality** based (local MEMORY > local NOF > first remote MEMORY > first remote NOF > LOCAL_DISK > DISK). Once locality fails, all remote MEMORY replicas are treated as equivalent, so the choice among them is effectively arbitrary (master return order).

This ignores two signals that the transfer layer already tracks per local NIC:

1. **NUMA / PCIe topology distance** — which local NIC reaches a given remote endpoint, and whether that NIC is same-NUMA or cross-NUMA relative to the destination GPU buffer.
2. **Real-time NIC load** — in-flight bytes and EWMA bandwidth per NIC.

On our 8x H20 nodes, cross-NUMA NIC selection shows up to ~2.4x tail-latency inflation versus same-NUMA, so picking an arbitrary remote replica can route a transfer through a congested or cross-NUMA path even when a better peer exists.

**Motivation**

The information needed to make a better choice is **already computed** inside TENT's `DeviceSelector` (`mooncake-transfer-engine/tent/src/transport/rdma/quota.cpp`). Its existing NIC-selection score is:

```cpp
// quota.cpp:129-138
// Score formula: predicted_time × numa_penalty + random_jitter
double predicted_time = (inflight_bytes + slice_bytes) / ewma_bandwidth;
double rank_penalty = numa_tier_weights[rank]; // [1.0, 5.0, 10.0] = same-NUMA / cross-NUMA / cross-socket
double score = predicted_time * rank_penalty;
```

Today this score is only used to pick the **local NIC** for an already-chosen transfer. The same score can rank **candidate remote replicas** at near-zero extra cost, since `inflight_bytes`, `ewma_bandwidth`, and the NUMA priority matrix are maintained continuously regardless.

**Proposed change**

Make remote-replica selection topology- and load-aware, reusing the existing `DeviceSelector` score (no new metric collection, no new scoring system):

1. When `SelectBestReplica` falls through to the remote-MEMORY branch with more than one candidate, score each candidate via the best local NIC's `predicted_time × numa_penalty` and pick the minimum, instead of returning `first_memory`.
2. Keep all existing fast paths unchanged: local-endpoint hit returns immediately; single-replica and the DISK/LOCAL_DISK ordering are untouched.
3. Fall back to the current behavior (first candidate) if scoring is unavailable, so the change is fully backward compatible.
4. Add a unit test covering the multi-remote-MEMORY case and a short note in the relevant doc.

This is a small, self-contained change (~20-30 LOC) and is fully backward compatible — single-replica and local-hit workloads behave exactly as before. I'd be happy to submit a PR if maintainers agree on the direction.

This also lines up with the "QoS-aware" scheduling direction in the roadmap (#1883); it's a minimal first step that only touches replica ranking, not the transfer path itself.

One question: is there any planned or in-flight work on replica selection / placement (e.g. dynamic replica management in #1100) that this should coordinate with, and would you prefer the ranking logic to live in `mooncake-store` (client side) or be pushed down into the TENT `DeviceSelector` API?

### Before submitting a new issue...

- [x] Make sure you already searched for relevant issues and read the documentation

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.