kvcache-ai / kvcache-ai/Mooncake

[RFC]: TENT-Native UB/URMA Transport

Open
#3,093 1 comment 0 reactions 0 assignees View on GitHub
RFC
Dominant language
C++
Stars
6.6k
Forks
1.2k
Avg merge
3d 5h
Merged PRs (30d)
312

Description

### Changes proposed

## Summary

This RFC proposes adding a **TENT-native UB transport** backed by the openEuler UMDK URMA APIs.

The goal is to make UB a first-class TENT transport whose runtime-facing interfaces and architectural responsibilities are aligned with the existing TENT RDMA transport.

The intended data path is:

```text
TENT Request
→ UbTask / UbSlice
→ TENT UB Workers
→ UbContext / UbEndpoint / UbJfc
→ URMA Adapter
→ liburma
```

TENT owns task scheduling, slicing, endpoint lifecycle, telemetry, cancellation, and failure recovery. The URMA adapter is limited to native resource management and URMA API translation.

The initial hardware scope is **host DRAM-to-DRAM transfer**. Additional memory types should be enabled only after their registration, topology, and access semantics have been validated.

---

## Related Work

UB/URMA support was previously introduced for the classic Mooncake Transfer Engine in #1773, #1805, and #1855. That work established the underlying URMA device, memory-registration, endpoint, build, mock, and testing capabilities.

This RFC builds on that foundation but focuses on a different architectural scope: implementing UB as a TENT-native transport with TENT-owned tasks, slices, workers, endpoint lifecycle, scheduling signals, cancellation, and failure recovery.

---

## Motivation

TENT provides transport-independent runtime capabilities such as:

- transport selection;

- task and slice lifecycle management;

- admission control and priority scheduling;

- cancellation;

- receiver-side credits;

- runtime bandwidth estimation;

- path failover and health monitoring.

To participate fully in this model, UB should expose the same logical capabilities as TENT RDMA.

A native implementation allows the TENT scheduler to observe UB runtime information such as:

- UB device availability;

- Jetty and remote-path health;

- JFC completion pressure;

- outstanding requests and inflight bytes;

- estimated bandwidth;

- endpoint generations;

- timeout and error statistics.

---

## Proposed Architecture

```mermaid
flowchart TD
APP[Application / Store] --> RT[TENT Runtime]

RT --> AQ[Admission Queue]
AQ --> TS[Transport Selector]
TS --> RC[Receiver Credit Gate]

RC --> UB[UbTransport]
UB --> WK[UbTask / UbSlice / Workers]

WK --> PS[Device and Path Selection]
PS --> CTX[UbContext]
PS --> ES[EndpointStore]

ES --> EP[UbEndpoint]
CTX --> JFC[UbJfc]

CTX --> UA[URMA Adapter]
EP --> UA
JFC --> UA
UA --> LIB[liburma / UMDK]

JFC --> TM[Telemetry and Health]
TM --> WK
TM --> AQ
```

The main architectural boundary is:

> TENT owns scheduling and policy. The URMA adapter owns only native URMA operations and resource lifetime.

---

## Alignment with TENT RDMA

The UB implementation should follow the same logical decomposition and public semantics as TENT RDMA.

| TENT RDMA | Proposed TENT UB |
| --- | --- |
| `RdmaTransport` | `UbTransport` |
| `RdmaSubBatch` | `UbSubBatch` |
| `RdmaTask` | `UbTask` |
| `RdmaSlice` | `UbSlice` |
| `RdmaContext` | `UbContext` |
| `RdmaCQ` | `UbJfc` |
| `RdmaEndPoint` | `UbEndpoint` |
| QP | Jetty |
| CQ | JFC |
| MR | URMA target segment |
| GID | EID |
| `ibv_post_send()` | `urma_post_jetty_send_wr()` |
| `ibv_poll_cq()` | `urma_poll_jfc()` |

The TENT-facing interfaces and state semantics should remain consistent, while native URMA terminology should be preserved inside the transport.

---

## TENT Transport Interface

`UbTransport` should directly implement the common TENT transport operations:

```text
install()
uninstall()

allocateSubBatch()
freeSubBatch()

submitTransferTasks()
getTransferStatus()

addMemoryBuffer()
removeMemoryBuffer()

getEstimatedBandwidth()

supportsCancellation()
cancelTransferTask()
```

Notification support may remain disabled initially until a native UB notification path is implemented.

TENT runtime code should not directly depend on URMA handles such as:

```text
urma_jetty_t
urma_jfc_t
urma_target_seg_t
```

---

## Proposed File Layout

```text
mooncake-transfer-engine/tent/
├── include/tent/transport/ub/
│ ├── ub_transport.h
│ ├── context.h
│ ├── jfc.h
│ ├── endpoint.h
│ ├── endpoint_store.h
│ ├── slice.h
│ ├── workers.h
│ ├── buffers.h
│ ├── rail_monitor.h
│ ├── quota.h
│ ├── params.h
│ └── urma_adapter.h

└── src/transport/ub/
├── CMakeLists.txt
├── ub_transport.cpp
├── context.cpp
├── jfc.cpp
├── endpoint.cpp
├── endpoint_store.cpp
├── workers.cpp
├── buffers.cpp
├── rail_monitor.cpp
├── quota.cpp
└── urma_adapter.cpp
```

The first implementation does not need to complete every optional scheduling feature, but the module boundaries should support later extensions without redesigning the data path.

---

## URMA Adapter Boundary

The URMA adapter should only expose low-level operations:

- URMA initialization and shutdown;

- device and EID discovery;

- context creation and destruction;

- asynchronous event FD exposure;

- local segment registration;

- remote segment import;

- JFC/JFCE creation and polling;

- Jetty creation, import, bind, reset, quiesce, and destruction;

- READ/WRITE work-request submission.

It should not contain:

- TENT task scheduling;

- slice spraying;

- receiver-credit policy;

- transport selection;

- retry policy;

- admission control;

- rail-health decisions.

Full URMA asynchronous-event consumption and acknowledgement may be added later if required by the provider and recovery model.

---

## Endpoint Lifecycle

At the conceptual level, UB endpoints should follow the same forward-only lifecycle semantics as TENT RDMA:

```mermaid
stateDiagram-v2
[*] --> UNINIT
UNINIT --> HANDSHAKING
HANDSHAKING --> READY
HANDSHAKING --> DESTROYING
READY --> DESTROYING
DESTROYING --> DESTROYED
DESTROYED --> [*]
```

The implementation may use additional internal states such as:

```text
PREPARED
BINDING
FAILED
```

Important requirements:

- lifecycle states only move forward;

- failed endpoint incarnations are retired rather than reconnected;

- each reconstruction creates a new generation;

- slices retain the endpoint generation used for submission;

- stale completions cannot update a newer endpoint;

- already-posted operations must be fenced or drained before native resources are destroyed;

- destruction and cleanup should be retryable when the provider reports a temporary busy or failure condition.

---

## Topology and Path Selection

UB should be represented as an independent TENT transport and topology type:

```text
TransportType::UB
Topology::NIC_UB
```

UB devices should not be represented as RDMA NICs.

Where available, topology information should include:

- native device name;

- stable TENT-facing device identity;

- device and EID indices;

- EID;

- NUMA node;

- PCI location;

- active state.

The initial path-selection model may consider:

```text
topology locality
local-device availability
remote-path availability
device and endpoint inflight pressure
transport quota
endpoint readiness
runtime bandwidth samples
retry history
```

This allows TENT to evaluate UB paths using the same logical model used for RDMA without exposing Jetty or JFC internals to the runtime.

---

## Native Control-Plane Bootstrap

UB should use a dedicated TENT bootstrap description rather than reusing RDMA QP fields.

The bootstrap message should include:

```text
protocol version
segment identity
local and peer NIC paths
local UB device identity
EID
Jetty IDs
endpoint generation
segment generation
capabilities
```

```mermaid
sequenceDiagram
participant W as UB Workers
participant S as EndpointStore
participant E as Local Endpoint
participant C as TENT Control Plane
participant P as Peer UB Transport
participant U as URMA Adapter

W->>S: getOrCreate(path)
S->>E: create endpoint generation N
E->>U: create local Jetty set
E->>C: bootstrap(EID, Jetty IDs, generation)
C->>P: invoke UB bootstrap handler
P->>U: create and bind peer Jetty set
P-->>C: peer EID and Jetty IDs
C-->>E: bootstrap response
E->>U: import and bind remote Jettys
E-->>S: endpoint READY
```

---

## Scheduling and Runtime Signals

The native implementation should expose sufficient runtime information for TENT scheduling:

```text
per-device inflight bytes
per-endpoint outstanding operations
completion latency
EWMA bandwidth
timeout count
completion error count
endpoint reconstruction count
local device state
remote path state
```

The initial implementation may use a deterministic rule-based score. More advanced QoS or deadline-aware arbitration can be added incrementally.

---

## Failure and Cancellation Semantics

The implementation should distinguish at least:

```text
local device error
remote path error
endpoint error
memory registration or import error
timeout
cancellation before post
unknown completion error
```

Suggested behavior:

- local-device failures retry through another local UB context;

- remote-path failures pause the path and try another remote device;

- endpoint failures retire the current generation;

- remote-segment generation changes trigger a new import;

- timeout retires the affected endpoint and initiates safe reconstruction;

- invalid address, permission, or range errors fail without retry.

Cancellation should follow the existing TENT best-effort model:

- unposted slices should not reach the device;

- queued and retrying slices may be canceled;

- already-posted operations must still be completed or safely fenced;

- buffers must not be reused while native DMA may still access them.

---

## Proposed Implementation Phases

### Phase 1: Native skeleton

- add `TransportType::UB`;

- add `Topology::NIC_UB`;

- add loader and CMake integration;

- add `UbTransport` and module skeletons;

- add an injectable URMA adapter.

### Phase 2: Context and memory

- discover URMA devices and EIDs;

- initialize contexts and JFC resources;

- register local segments;

- import remote segments;

- publish TENT-native UB metadata.

### Phase 3: Endpoint bootstrap

- add the UB control-plane message;

- implement Jetty bootstrap;

- implement `EndpointStore`;

- implement endpoint lifecycle and generations.

### Phase 4: READ/WRITE data path

- implement UB tasks and slices;

- implement posting and polling workers;

- post URMA READ/WRITE operations;

- dispatch JFC completions;

- expose task status.

### Phase 5: Reliability and scheduling

- cancellation;

- error classification;

- path retry and failover;

- timeout and endpoint reconstruction;

- runtime bandwidth estimation;

- admission-queue and receiver-credit integration.

Each phase should remain independently buildable, testable, and reviewable.

---

## Validation

CI should use an injectable/mock URMA adapter to validate deterministic behavior without UB hardware.

Real URMA hardware validation remains a separate acceptance requirement before the transport can be considered production-ready.

Hardware validation should cover:

- two-node READ and WRITE correctness;

- mixed-operation consistency checking;

- multiple UB devices and paths;

- endpoint reuse and generation isolation;

- peer restart;

- remote segment re-registration;

- UB port down/up;

- timeout and cancellation;

- safe repeated shutdown;

- performance comparison against the existing UB implementation.

The architecture is designed to support multi-path scheduling and recovery, but those capabilities should not be considered validated until the real-hardware test matrix passes.

---

## Non-Goals

The initial work does not need to:

- replace existing transport implementations;

- support GPU or other non-DRAM memory immediately;

- implement every TENT QoS capability;

- implement native UB notifications in the first version;

- rename URMA concepts using RDMA terminology;

- place scheduling logic inside the URMA adapter;

- redesign unrelated TENT runtime APIs.

Feedback on the module boundaries, bootstrap protocol, metadata format, PR decomposition, and hardware validation plan would be appreciated

### Before submitting a new issue...

- [x] Make sure you already searched for relevant issues and read the [documentation](https://kvcache-ai.github.io/Mooncake/)

Contributor guide

Open the contributing guide

Research direction

Start by reading the existing TENT RDMA implementation and comparing it with the proposed files under tent/include/tent/transport/ub/ and tent/src/transport/ub/. Trace the TENT transport interface, CMake integration, and injectable URMA adapter boundary before choosing a phase. Done should be a separately buildable, testable phase with deterministic mock-URMA coverage; hardware validation is a separate requirement.

Written by the indexing model from the issue text.

Assessment

Tech stack
cmake, cpp
Domain
build-system, distributed-systems, networking
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Quiet
Clarity
Needs clarification
Newbie friendliness
25/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.