[TE] Enabling UB Transport on the Kunpeng SuperNode Phase 1 by zchuango · Pull Request #1805 · kvcache-ai/Mooncake

zchuango · 2026-04-02T09:29:11Z

Description

Overview

This PR introduces UbTransport, a new transport implementation for Mooncake's Transfer Engine that enables high-performance data transfer on Kunpeng 950 platforms using the UB (Unified Bus) protocol and URMA (Unified Remote Memory Access) APIs from the openEuler UMDK.

Related Issue: #1773

Motivation

Hardware: Kunpeng 950 CPUs with native UB fabric support
Workload: 8MB-scale embedding tables and KV cache transfers for generative recommendation
Target: <10ms latency for 8MB transfer, 400Gbps bandwidth saturation
Protocol: URMA provides RDMA-like memory semantics optimized for Kunpeng's on-die interconnect

Build & Compile

Prerequisites

OpenEuler 24.03 (LTS-SP3) download
UMDK installed: yum install umdk-urma-devel or build from source

CMake Configuration

# Clone Mooncake
git clone https://github.com/kvcache-ai/Mooncake.git
cd Mooncake

# Configure with UB transport enabled
mkdir build && cd build
cmake .. -DUSE_UB=ON \
         -DURMA_INCLUDE_DIR=/usr/include \
         -DURMA_LIBRARY=/usr/lib64/liburma.so

# Build
make -j$(nproc)

Verification

# Check UB transport is registered
./mooncake_server --list-transports
# Expected output: rdma, tcp, nvlink, ub

Run & Test

Single-node benchmark

# Terminal 1: Target
./transfer_engine_bench \
    --mode=target \
    --protocol=ub \
    --device_name=urma0 \
    --local_server_name=127.0.0.1 \
    --metadata_server=P2PHANDSHAKE

# Terminal 2: Initiator
./transfer_engine_bench \
    --mode=initiator \
    --protocol=ub \
    --device_name=urma0 \
    --metadata_server=P2PHANDSHAKE \
    --segment_size=8388608 \
    --batch_size=1\
    --segment_id=127.0.0.1:$PORT

Multi-device benchmark

# Discover multiple URMA devices
./transfer_engine_bench \
    --protocol=ub \
    --device_name=urma0,urma1,urma2,urma3

Phase Contribution Plan

Following suggestion, we propose four phased PRs for incremental review and integration:

Phase 1: Enable UbTransport + UrmaEndpoint ⬅️ Current PR

Review focus: Transport interface compliance, WorkerPool design, error handling

mooncake-common/
   ├── common.cmake    # add USE_UB option
mooncake-transfer-engine/
   ├── include/
   │   └── mooncake/
   │       └── kunpeng_transport/
   │           ├── ub_transport.h     # Transport interface implementation
   │           ├── ub_context.h       # UbDevice Init and TaskPool scheduling And event polling
   │           ├── ub_endpoint.h      # ub device endpoint interface will imp
   │           └── urma_endpoint.h    # RMA Read/Write operations
   ├── src/
   │   └── kunpeng_transport/
   │       ├── CMakeLists.txt
   │       ├── ub_context.cpp
   │       ├── ub_transport.cpp
   │       └──  urma_endpoint.cpp
   ├── CMakeLists.txt     # USE_UB option integration ub_transport
   └── tests/
         └── ub_transport_test.cpp

Phase 2: CI/Testing Infrastructure With Kunpeng UB Hardware Env
Phase 3: Enable UbTransport + OBMM Feature
Phase 4: Next TE Tent Module Integration with UbTransport

Acknowledgments

Thanks to @stmatengss and the Mooncake team for the phased contribution guidance.
Ready for review: Phase 1 implementation is complete and tested. Please advise on any architectural adjustments before we proceed with detailed code review.

@stmatengss @alogfans @ascend-direct-dev

Module

Type of Change

How Has This Been Tested?

Checklist

I have performed a self-review of my own code.
I have formatted my own code using ./scripts/code_format.sh before submitting.
I have updated the documentation.
I have added tests to prove my changes are effective.

…ort_dev # Conflicts: # mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp # mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp

gemini-code-assist

Code Review

This pull request introduces the 'ub' (User-space Bridge) transport protocol, specifically integrating the URMA (User-space RDMA) implementation. The changes include new configuration parameters, metadata structures, and a complete transport implementation within the kunpeng_transport directory. The topology discovery mechanism has been updated to support UB devices, and example benchmarks now include 'ub' protocol support. Review feedback highlights several critical memory management issues, such as incorrect use of delete on std::shared_ptr and malloc-allocated memory, as well as the need to return Slice objects to their cache instead of deleting them. Logic errors were also identified where getters returning vectors by value were incorrectly used for modifications, and an off-by-one error was found in a bounds check.

codecov-commenter · 2026-04-02T10:16:34Z

⚠️ Please install the to ensure uploads and comments are reliably processed by Codecov.

Codecov Report

❌ Patch coverage is 3.22581% with 60 lines in your changes missing coverage. Please review.

Files with missing lines	Patch %	Lines
mooncake-transfer-engine/src/transfer_metadata.cpp	4.16%	46 Missing ⚠️
...gine/example/transfer_engine_bench_with_notify.cpp	0.00%	6 Missing ⚠️
...-transfer-engine/example/transfer_engine_bench.cpp	0.00%	5 Missing ⚠️
mooncake-transfer-engine/src/config.cpp	0.00%	3 Missing ⚠️

📢 Thoughts on this report? Let us know!

…ort_dev

Copilot

Pull request overview

This PR adds an initial UB (Unified Bus) transport implementation to Mooncake Transfer Engine, targeting Kunpeng platforms via the openEuler UMDK URMA APIs. It introduces new transport/context/endpoint codepaths, integrates UB into build + transport registration, and extends metadata/topology to support UB device discovery and connection handshakes.

Changes:

Introduce UbTransport with URMA-backed UrmaContext / UrmaEndpoint, plus worker pool + endpoint store.
Add USE_UB build option and wire UB transport into TE init, multi-transport registry, metadata encoding/handshake, and topology discovery.
Add a UB “test” executable and update example benchmarks to accept --protocol=ub.

Reviewed changes

Copilot reviewed 24 out of 24 changed files in this pull request and generated 17 comments.

Show a summary per file

File	Description
mooncake-transfer-engine/tests/ub_transport_test.cpp	Adds a UB integration-style executable (currently behaves like a manual 2-process runner).
mooncake-transfer-engine/tests/CMakeLists.txt	Registers `ub_transport_test` in build + CTest under `USE_UB`.
mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp	Implements URMA context + endpoint (registration/import, jetty setup, posting/polling).
mooncake-transfer-engine/src/transport/kunpeng_transport/ub_transport.cpp	Implements the `Transport` interface for UB and wires it to contexts + handshake.
mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp	Implements endpoint store + worker pool (post-send scheduling, polling, redispatch).
mooncake-transfer-engine/src/transport/kunpeng_transport/CMakeLists.txt	Adds CMake target for UB transport sources + URMA linkage.
mooncake-transfer-engine/src/transport/CMakeLists.txt	Adds `kunpeng_transport` subdir and objects into `transport` under `USE_UB`.
mooncake-transfer-engine/src/transfer_metadata.cpp	Extends handshake + segment descriptor (encode/decode) for UB (EID, jetty_num, tseg).
mooncake-transfer-engine/src/transfer_engine_impl.cpp	Auto-installs UB transport when `USE_UB` is enabled.
mooncake-transfer-engine/src/topology.cpp	Adds UB device discovery via URMA + sysfs and switches CPU topology discovery under `USE_UB`.
mooncake-transfer-engine/src/multi_transport.cpp	Registers `"ub"` protocol in `MultiTransport`.
mooncake-transfer-engine/src/config.cpp	Extends env parsing to also populate UB JFC/JFCE config from `MC_NUM_CQ_PER_CTX`.
mooncake-transfer-engine/src/CMakeLists.txt	Links UB transport target when `USE_UB` is enabled.
mooncake-transfer-engine/include/transport/transport.h	Extends `Transport::Slice` union with UB-specific fields.
mooncake-transfer-engine/include/transport/kunpeng_transport/urma_endpoint.h	Adds URMA context/endpoint declarations and URMA-specific helpers.
mooncake-transfer-engine/include/transport/kunpeng_transport/ub_transport.h	Adds UB transport interface declaration.
mooncake-transfer-engine/include/transport/kunpeng_transport/ub_endpoint.h	Adds abstract UB endpoint interface.
mooncake-transfer-engine/include/transport/kunpeng_transport/ub_context.h	Adds UB context abstraction and worker pool / endpoint store declarations.
mooncake-transfer-engine/include/transfer_metadata.h	Extends metadata structures for UB (EID, tseg/l_seg_index, jetty_num).
mooncake-transfer-engine/include/config.h	Adds UB-related config fields (jfc/jfce counts, eid_index, max_seg_size, etc.).
mooncake-transfer-engine/example/transfer_engine_bench.cpp	Adds UB transport install path and changes host free path for `--protocol=ub`.
mooncake-transfer-engine/example/transfer_engine_bench_with_retry.cpp	Adds UB transport install path.
mooncake-transfer-engine/example/transfer_engine_bench_with_notify.cpp	Adds UB transport install path.
mooncake-common/common.cmake	Adds `USE_UB` option and compile definition.

Comments suppressed due to low confidence (1)

mooncake-transfer-engine/src/topology.cpp:490

When USE_UB is enabled, all_hca is a std::vector<UBDevice>, but discoverCudaTopology is declared to take const std::vector<InfinibandDevice>&. This call will not compile in CUDA/MUSA/HIP builds. Consider templating/overloading discoverCudaTopology for UBDevice, or using a shared device struct type for both RDMA and UB discovery paths.

#if defined(USE_CUDA) || defined(USE_MUSA) || defined(USE_HIP)
    for (auto &ent : discoverCudaTopology(all_hca)) {
        matrix_[ent.name] = ent;
    }

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

alogfans · 2026-04-03T07:41:24Z

The usage guide should be included in the docs. In addition, what's the main advantage of using this approach, rather than other compatible xports?

zchuango · 2026-04-03T10:51:52Z

The usage guide should be included in the docs. In addition, what's the main advantage of using this approach, rather than other compatible xports?

Sure, I will add some more details about the usage guide later. By the way, which directory should I write it to? @alogfans

Here are some of my thoughts on the advantages of the current proposal:

It will improve and expand the current mainstream protocol framework of the TE module, enabling the mooncake TE to run on more communication protocols and related hardware, thus expanding the mooncake's application platform.
The UB protocol is a peer-to-peer communication protocol comparable to NVlink, CXL, and RDMA. Currently, there is already hardware support on the Kunpeng 950 super node. Based on the UB protocol, we can implement a URMA component that supports DMA semantics, as well as an OBMM component that supports LD/ST semantics. As the technology evolves, these components can support core requirements for lower latency and higher bandwidth in model inference.

@alogfans @stmatengss welcome any suggestions for this issue. I am very happy to participate in the construction of the community ecosystem together.

stmatengss · 2026-04-05T04:00:50Z

Please fix this CI error.

https://github.com/kvcache-ai/Mooncake/actions/runs/23950310040/job/69858773912?pr=1805

zchuango · 2026-04-05T05:14:38Z

Please fix this CI error.

https://github.com/kvcache-ai/Mooncake/actions/runs/23950310040/job/69858773912?pr=1805

@stmatengss could you please rerun it again? I haven't been able to pinpoint the problem yet. I tried running CI/Testing on my personal repository pull request zchuango#2 and it can run normally, so I suspect it might be an intermittent issue.

stmatengss · 2026-04-06T07:46:23Z

Please update the documentation at https://github.com/kvcache-ai/Mooncake/tree/main/docs/source/design/transfer-engine or consider it as the next step.

stmatengss · 2026-04-06T07:51:56Z

@alogfans Please double-check

zchuango · 2026-04-07T01:43:30Z

Please update the documentation at https://github.com/kvcache-ai/Mooncake/tree/main/docs/source/design/transfer-engine or consider it as the next step.

Yes, I plan to update the documentation in the next step.

…ai#1805) * add the support for ub transport * refine the code format * change some log for ub transport * change ub_transport to kunpeng_transport * add the support for ub transport * refine the code format * change some log for ub transport * change ub_transport to kunpeng_transport * clean some comments and unused code * fix the config.h merge error * fix the code with review suggestion * add the usage docs and refine some code for code review * fix some comments with chinese and modify the CMakeLists.txt file

zchuango added 12 commits March 19, 2026 21:50

add the support for ub transport

5966dab

refine the code format

21e9633

change some log for ub transport

7d8dfb4

change ub_transport to kunpeng_transport

7650d7d

add the support for ub transport

7914cfc

refine the code format

3da4eda

change some log for ub transport

147bc77

change ub_transport to kunpeng_transport

a1fea3c

clean some comments and unused code

d4b2235

Merge branch 'main' into ub_transport_dev

ee81198

Merge remote-tracking branch 'origin/ub_transport_dev' into ub_transp…

1424e85

…ort_dev # Conflicts: # mooncake-transfer-engine/src/transport/kunpeng_transport/ub_context.cpp # mooncake-transfer-engine/src/transport/kunpeng_transport/urma_endpoint.cpp

fix the config.h merge error

48f9099

zchuango requested review from alogfans, chestnut-Q and doujiang24 as code owners April 2, 2026 09:29

gemini-code-assist Bot reviewed Apr 2, 2026

View reviewed changes

github-actions Bot added run-ci Transfer Engine labels Apr 2, 2026

zchuango added 2 commits April 2, 2026 18:23

fix the code with review suggestion

6549c11

Merge remote-tracking branch 'origin/ub_transport_dev' into ub_transp…

1fb6cf8

…ort_dev

stmatengss requested a review from Copilot April 2, 2026 11:40

Copilot started reviewing on behalf of stmatengss April 2, 2026 11:41 View session

Copilot AI reviewed Apr 2, 2026

View reviewed changes

zchuango requested review from ShangmingCai, stmatengss and ykwd as code owners April 3, 2026 12:09

add the usage docs and refine some code for code review

76c8e46

zchuango force-pushed the ub_transport_dev branch from 9a61ac0 to 76c8e46 Compare April 3, 2026 14:48

stmatengss reviewed Apr 5, 2026

View reviewed changes

Comment thread mooncake-transfer-engine/include/config.h Outdated

stmatengss reviewed Apr 5, 2026

View reviewed changes

Comment thread mooncake-transfer-engine/include/transfer_metadata.h

stmatengss reviewed Apr 5, 2026

View reviewed changes

Comment thread mooncake-transfer-engine/src/config.cpp

fix some comments with chinese and modify the CMakeLists.txt file

02ef725

stmatengss reviewed Apr 6, 2026

View reviewed changes

Comment thread mooncake-transfer-engine/src/topology.cpp

stmatengss approved these changes Apr 6, 2026

View reviewed changes

alogfans merged commit 6da2572 into kvcache-ai:main Apr 7, 2026
17 checks passed

zchuango mentioned this pull request Apr 9, 2026

[TE] Enabling UB Transport on the Kunpeng SuperNode Phase 2 #1855

Merged

25 tasks

This was referenced Jul 3, 2026

[Transfer Engine] Enable UB Transport in TENT for Kunpeng SuperNode Phase 3 Le1zyCatt/Mooncake#1

Open

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

Open

zchuango mentioned this pull request Jul 10, 2026

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

Closed

27 tasks

zchuango mentioned this pull request Jul 24, 2026

[RFC]: TENT-Native UB/URMA Transport #3093

Open

1 task

Conversation

zchuango commented Apr 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Overview

Motivation

Build & Compile

Run & Test

Phase Contribution Plan

Acknowledgments

Module

Type of Change

How Has This Been Tested?

Checklist

Uh oh!

gemini-code-assist Bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

codecov-commenter commented Apr 2, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Copilot AI left a comment

Choose a reason for hiding this comment

Pull request overview

Reviewed changes

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

alogfans commented Apr 3, 2026

Uh oh!

zchuango commented Apr 3, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

stmatengss commented Apr 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

zchuango commented Apr 5, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

stmatengss commented Apr 6, 2026

Uh oh!

Uh oh!

stmatengss commented Apr 6, 2026

Uh oh!

zchuango commented Apr 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

5 participants

zchuango commented Apr 2, 2026 •

edited

Loading

codecov-commenter commented Apr 2, 2026 •

edited

Loading

zchuango commented Apr 3, 2026 •

edited

Loading

stmatengss commented Apr 5, 2026 •

edited

Loading

zchuango commented Apr 5, 2026 •

edited

Loading