kvcache-ai / kvcache-ai/Mooncake
[RFC]: TENT-Native UB/URMA Transport
- 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
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