[TE/TENT] Enable UB Transport in TENT on Kunpeng SuperNode (Phase 3)

[zchuango]

Description

Overview

This PR continues the phased UB transport work for Mooncake Transfer Engine on the Kunpeng SuperNode.

Phase 1 introduced the initial UbTransport and URMA endpoint implementation for the legacy Transfer Engine. Phase 2 added mock URMA support, unit tests, and CI/testing infrastructure for UB transport. This Phase 3 PR enables UB transport in TENT by adding a TENT-side UB transport adapter and the required control-plane integration.

This PR is submitted early to make the UB-in-TENT integration visible for incremental review. The current scope is to make UB recognizable, selectable, bootstrappable, and usable from TENT while reusing the existing legacy UbTransport data path. Follow-up commits may further refine the memory allocation/alignment path, metadata abstraction, and deeper TENT-native integration.

Related work:

• [RFC]: Enable UB transport via Kunpeng UMDK URMA apis kvcache-ai/Mooncake#1773
• [TE] Enabling UB Transport on the Kunpeng SuperNode Phase 1 kvcache-ai/Mooncake#1805
• [TE] Enabling UB Transport on the Kunpeng SuperNode Phase 2 kvcache-ai/Mooncake#1855
• [TENT] Enhanced QoS and Slice Spraying for TENT kvcache-ai/Mooncake#2048

---

Motivation

UB transport support already exists in the legacy Transfer Engine path, but TENT currently cannot select or drive UB as a first-class transport.

This prevents Kunpeng SuperNode UB hardware from being used through the newer TENT transport abstraction. To close this gap, this PR introduces a compatibility layer that maps TENT transport APIs, segment metadata, memory registration, transfer batches, bootstrap information, and status queries to the existing UB transport implementation.

The goal of this phase is not to rewrite the full URMA/UB data path. Instead, this PR provides a practical bridge so that UB can be used through TENT first, while leaving room for a more native TENT UB implementation later.

---

Major Changes

1. Add UB as a TENT transport type

This PR adds UB support to the TENT transport type system, including:

• TransportType::UB
• "ub" string parsing in the TENT selector
• UB transport creation in the TENT transport loader
• UB transport type exposure through pybind
• CMake integration for the TENT UB transport target

This allows TENT to recognize UB as a selectable transport type.

---

2. Add UbTentTransport

This PR introduces UbTentTransport, a TENT transport adapter for the existing legacy UbTransport.

UbTentTransport implements the TENT transport interface and delegates UB data path operations to the legacy UbTransport.

It handles:

• TENT transport installation and uninstallation
• UB device selection
• Local memory registration and unregistration
• TENT buffer descriptor conversion
• TENT batch allocation and release
• TENT transfer request conversion
• Transfer status query and status mapping

This keeps the existing URMA/UB data path mostly unchanged and limits this PR to the TENT integration layer.

---

3. Add UbTentMetadataBridge

Legacy UbTransport depends on the old TransferMetadata interface, while TENT uses SegmentManager and TENT-style segment descriptors.

This PR adds UbTentMetadataBridge to bridge the two metadata systems. The bridge is responsible for:

• Looking up local and remote TENT segments
• Converting TENT segment descriptors to legacy Transfer Engine segment descriptors
• Extracting UB-specific transport attributes such as EID and remote segment information
• Providing the metadata interface expected by legacy UbTransport

This allows the existing UB transport logic to work under the TENT segment management model.

---

4. Add UB bootstrap support to the TENT control plane

UB requires endpoint bootstrap information exchange before data transfer.

This PR adds a UB bootstrap path to the TENT control plane, including:

• RpcFuncID::BootstrapUb
• UB bootstrap descriptor structure
• ControlClient::bootstrapUb()
• ControlService::onBootstrapUb()
• Callback registration support for UB bootstrap handling

With this change, UB endpoint bootstrap can be performed through the TENT control-plane RPC path.

---

5. Fix UB device selection and device_name propagation

This PR fixes UB device selection in the TENT path.

Before this change, device_name propagation was incomplete for UB, which could cause UB to discover or select the wrong device even when the expected UB device was provided by configuration or benchmark parameters.

This PR ensures that UB can use the configured device name, such as bonding_dev_0, through the TENT configuration path.

The intended selection order is:

1. transports/ub/device_name
2. MC_UB_DEVICE_NAME
3. Automatic UB device discovery

This makes dual-node UB/TENT testing and production deployment more deterministic.

---

6. Improve URMA memory registration and cleanup behavior

During UB/TENT integration, several URMA resource lifetime issues were found and fixed.

This PR improves:

• Page-aligned memory registration for URMA
• Avoiding repeated registration of the same host virtual address across multiple URMA contexts
• Reusing the primary registered segment where possible
• Releasing local segment references before unregistering memory
• Avoiding incorrect urma_uninit() usage during per-context teardown
• Safer cleanup paths when initialization partially fails

These changes make UB transport more stable in real URMA environments.

---

7. Add tests and documentation

This PR adds TENT UB related tests and documentation, covering:

• TENT UB transport install/uninstall path
• Memory buffer registration and removal
• Sub-batch allocation and release
• Mock transfer submit and status query
• Dual-node UB/TENT integration test entry
• Phase 3 UB/TENT test guide

The dual-node test path is documented in:

mooncake-transfer-engine/tent/docs/ub_tent_transport_guide.md

---

Scope

This PR focuses on enabling UB transport inside TENT by adapting the existing UB implementation.

This PR does not redesign the full UB data path and does not introduce a full TENT-native UB backend. It also does not change the existing TENT RDMA slice-spraying implementation.

The current implementation should be understood as a Phase 3 compatibility/integration layer:

• TENT can select UB.
• TENT can bootstrap UB endpoints.
• TENT can register UB memory.
• TENT can submit transfer requests through UB.
• The actual UB data path still reuses legacy UbTransport.

---

Current Limitations and Follow-up Work

Memory alignment and allocator integration

The current implementation fixes URMA registration stability with page-aligned memory handling and primary-context registration/adoption.

As follow-up work, we will evaluate whether this page-alignment logic should be moved to or replaced by a common allocator. This would avoid keeping transport-specific allocation/alignment logic scattered inside the UB/TENT adapter path.

---

TENT slice spraying integration

TENT’s design intends slice spraying to be a transport-agnostic scheduling capability at the TENT layer. The current upstream implementation mainly provides RDMA-oriented QoS and slice spraying.

This PR does not yet integrate UB into the TENT-native slice scheduling path. UB is currently exposed as a selectable TENT transport while reusing the legacy UB data path.

A deeper integration with TENT slice scheduling/QoS can be evaluated in follow-up work once the UB backend is refactored away from the legacy UbTransport adapter.

---

Metadata abstraction cleanup

This PR makes selected TransferMetadata methods virtual so that UbTentMetadataBridge can provide the metadata and bootstrap behavior expected by legacy UbTransport.

This is a practical bridge for Phase 3, but it is not necessarily the cleanest long-term abstraction. A future refactor may introduce a narrower UB metadata/handshake interface or decouple UbTransport from the legacy TransferMetadata dependency.

---

Native TENT UB backend

This PR is not a full TENT-native UB rewrite.

A native TENT UB backend, including deeper integration with TENT scheduling, QoS, and slice spraying, will be considered based on the evolution and stabilization of TENT transport/scheduler interfaces.

---

Module

• Transfer Engine (mooncake-transfer-engine)
• Mooncake Store (mooncake-store)
• Mooncake EP (mooncake-ep)
• Mooncake PG (mooncake-pg)
• Integration (mooncake-integration)
• P2P Store (mooncake-p2p-store)
• Python Wheel (mooncake-wheel)
• Common (mooncake-common)
• Mooncake RL (mooncake-rl)
• CI/CD
• Docs
• Other

---

Type of Change

• Bug fix
• New feature
• Refactor
• Breaking change
• Documentation update
• Performance improvement
• Other

---

How Has This Been Tested?

Build with UB and TENT enabled

mkdir -p build
cd build

cmake .. \
  -DUSE_UB=ON \
  -DUSE_TENT=ON \
  -DCMAKE_BUILD_TYPE=RelWithDebInfo

cmake --build . -j$(nproc)

Run TENT UB unit tests

ctest -R ub_tent_transport_test --output-on-failure

Run UB transport tests

ctest -R ub_transport_test --output-on-failure

Dual-node UB/TENT validation

This PR has been validated on a dual-node Kunpeng UB setup with real UB devices.

The validation covers:

• TENT-side UB transport selection
• UB device_name propagation
• UB endpoint bootstrap through TENT control-plane RPC
• TENT segment publishing and lookup
• URMA memory registration
• End-to-end UB transfer through the TENT adapter path

The dual-node test procedure is documented in:

mooncake-transfer-engine/tent/docs/ub_tent_transport_guide.md

Test results

• Unit tests pass
• UB transport tests pass
• Dual-node UB/TENT integration test passes
• Manual validation on real Kunpeng UB hardware completed

---

Checklist

• I have performed a self-review of my own code
• I have formatted my code using ./scripts/code_format.sh
• I have run pre-commit run --all-files and all hooks pass
• I have updated the documentation
• I have added tests to prove my changes are effective
• For changes >500 LOC: I have filed or referenced related RFC / phased development issues or PRs

---

AI Assistance Disclosure

• No AI tools were used
• AI tools were used

AI tools were used to help draft and polish the PR description. The implementation, testing, and final verification are the responsibility of the human submitter.